Add tests for model_artifact update validations (opendatahub-io#284)

* Add tests for model_artifact update validations

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: dbasunag

tests/model_registry/exceptions.py

@@ -0,0 +1,10 @@
+class ModelRegistryResourceNotCreated(Exception):
+    pass
+
+
+class ModelRegistryResourceNotFoundError(Exception):
+    pass
+
+
+class ModelRegistryResourceNotUpdated(Exception):
+    pass

tests/model_registry/rest_api/__init__.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+import pytest
+
+from tests.model_registry.rest_api.constants import MODEL_REGISTRY_BASE_URI
+from tests.model_registry.rest_api.utils import register_model_rest_api, execute_model_registry_patch_command
+from utilities.constants import Protocols
+
+
+@pytest.fixture(scope="class")
+def model_registry_rest_url(model_registry_instance_rest_endpoint: str) -> str:
+    # address and port need to be split in the client instantiation
+    return f"{Protocols.HTTPS}://{model_registry_instance_rest_endpoint}"
+
+
+@pytest.fixture(scope="class")
+def model_registry_rest_headers(current_client_token: str) -> dict[str, str]:
+    return {
+        "Authorization": f"Bearer {current_client_token}",
+        "accept": "application/json",
+        "Content-Type": "application/json",
+    }
+
+
+@pytest.fixture(scope="class")
+def registered_model_rest_api(
+    request: pytest.FixtureRequest, model_registry_rest_url: str, model_registry_rest_headers: dict[str, str]
+) -> dict[str, Any]:
+    return register_model_rest_api(
+        model_registry_rest_url=model_registry_rest_url,
+        model_registry_rest_headers=model_registry_rest_headers,
+        data_dict=request.param,
+    )
+
+
+@pytest.fixture(scope="class")
+def updated_model_artifact(
+    request: pytest.FixtureRequest,
+    model_registry_rest_url: str,
+    model_registry_rest_headers: dict[str, str],
+    registered_model_rest_api: dict[str, Any],
+) -> dict[str, Any]:
+    model_artifact_id = registered_model_rest_api["model_artifact"]["id"]
+    assert model_artifact_id, f"Model artifact id not found: {registered_model_rest_api['model_artifact']}"
+    return execute_model_registry_patch_command(
+        url=f"{model_registry_rest_url}{MODEL_REGISTRY_BASE_URI}model_artifacts/{model_artifact_id}",
+        headers=model_registry_rest_headers,
+        data_json=request.param,
+    )

tests/model_registry/rest_api/conftest.py

@@ -0,0 +1,31 @@
+from typing import Any
+
+from utilities.constants import ModelFormat
+
+MODEL_REGISTER: dict[str, str] = {
+    "name": "model-rest-api",
+    "description": "Model created via rest call",
+    "owner": "opendatahub-tests",
+}
+MODEL_VERSION: dict[str, Any] = {
+    "name": "v0.0.1",
+    "state": "LIVE",
+    "author": "opendatahub-tests",
+    "description": "Model version created via rest call",
+}
+
+MODEL_ARTIFACT: dict[str, Any] = {
+    "name": "model-artifact-rest-api",
+    "description": "Model artifact created via rest call",
+    "uri": "https://huggingface.co/openai-community/gpt2/resolve/main/onnx/decoder_model.onnx",
+    "state": "UNKNOWN",
+    "modelFormatName": ModelFormat.ONNX,
+    "modelFormatVersion": "v1",
+    "artifactType": "model-artifact",
+}
+MODEL_REGISTER_DATA = {
+    "register_model_data": MODEL_REGISTER,
+    "model_version_data": MODEL_VERSION,
+    "model_artifact_data": MODEL_ARTIFACT,
+}
+MODEL_REGISTRY_BASE_URI = "/api/model_registry/v1alpha3/"

tests/model_registry/rest_api/constants.py

@@ -0,0 +1,104 @@
+from typing import Self, Any
+import pytest
+from tests.model_registry.rest_api.constants import MODEL_REGISTER, MODEL_ARTIFACT, MODEL_VERSION, MODEL_REGISTER_DATA
+from utilities.constants import DscComponents
+from tests.model_registry.constants import MR_NAMESPACE
+from simple_logger.logger import get_logger
+
+LOGGER = get_logger(name=__name__)
+
+
+@pytest.mark.parametrize(
+    "updated_dsc_component_state_scope_class, registered_model_rest_api",
+    [
+        pytest.param(
+            {
+                "component_patch": {
+                    DscComponents.MODELREGISTRY: {
+                        "managementState": DscComponents.ManagementState.MANAGED,
+                        "registriesNamespace": MR_NAMESPACE,
+                    },
+                },
+            },
+            MODEL_REGISTER_DATA,
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.usefixtures("updated_dsc_component_state_scope_class", "registered_model_rest_api")
+class TestModelRegistryCreationRest:
+    """
+    Tests the creation of a model registry. If the component is set to 'Removed' it will be switched to 'Managed'
+    for the duration of this test module.
+    """
+
+    @pytest.mark.parametrize(
+        "expected_params, data_key",
+        [
+            pytest.param(
+                MODEL_REGISTER,
+                "register_model",
+                id="test_validate_registered_model",
+            ),
+            pytest.param(
+                MODEL_VERSION,
+                "model_version",
+                id="test_validate_model_version",
+            ),
+            pytest.param(
+                MODEL_ARTIFACT,
+                "model_artifact",
+                id="test_validate_model_artifact",
+            ),
+        ],
+    )
+    def test_validate_model_registry_resource(
+        self: Self,
+        registered_model_rest_api: dict[str, Any],
+        expected_params: dict[str, str],
+        data_key: str,
+    ):
+        errors: list[str:Any]
+        created_resource_data = registered_model_rest_api[data_key]
+        if errors := [
+            {key: [expected_params[key], created_resource_data.get(key)]}
+            for key in expected_params.keys()
+            if not created_resource_data.get(key) or created_resource_data[key] != expected_params[key]
+        ]:
+            pytest.fail(f"Model did not get created with expected values: {errors}")
+        LOGGER.info(f"Successfully validated: {created_resource_data['name']}")
+
+    @pytest.mark.parametrize(
+        "updated_model_artifact, expected_param",
+        [
+            pytest.param(
+                {"description": "updated description"},
+                {"description": "updated description"},
+                id="test_validate_updated_artifact_description",
+            ),
+            pytest.param(
+                {"modelFormatName": "tensorflow"},
+                {"modelFormatName": "tensorflow"},
+                id="test_validate_updated_artifact_model_format_name",
+            ),
+            pytest.param(
+                {"modelFormatVersion": "v2"},
+                {"modelFormatVersion": "v2"},
+                id="test_validate_updated_artifact_model_format_version",
+            ),
+        ],
+        indirect=["updated_model_artifact"],
+    )
+    def test_create_update_model_artifact(
+        self,
+        updated_model_artifact: dict[str, Any],
+        expected_param: dict[str, Any],
+    ):
+        errors: list[dict[str, list[Any]]]
+        if errors := [
+            {key: [expected_param[key], updated_model_artifact.get(key)]}
+            for key in expected_param.keys()
+            if not updated_model_artifact.get(key) or updated_model_artifact[key] != expected_param[key]
+        ]:
+            pytest.fail(f"Model did not get updated with expected values: {errors}")
+        LOGGER.info(f"Successfully validated: {updated_model_artifact['name']}")

tests/model_registry/rest_api/test_model_registry_rest_api.py

@@ -0,0 +1,89 @@
+from typing import Any
+import requests
+import json
+from simple_logger.logger import get_logger
+from tests.model_registry.exceptions import (
+    ModelRegistryResourceNotCreated,
+    ModelRegistryResourceNotFoundError,
+    ModelRegistryResourceNotUpdated,
+)
+from tests.model_registry.rest_api.constants import MODEL_REGISTRY_BASE_URI
+
+LOGGER = get_logger(name=__name__)
+
+
+def execute_model_registry_patch_command(
+    url: str, headers: dict[str, str], data_json: dict[str, Any]
+) -> dict[Any, Any]:
+    resp = requests.patch(url=url, json=data_json, headers=headers, verify=False, timeout=60)
+    LOGGER.info(f"url: {url}, status code: {resp.status_code}, rep: {resp.text}")
+
+    if resp.status_code != 200:
+        raise ModelRegistryResourceNotUpdated(
+            f"Failed to update ModelRegistry resource: {url}, {resp.status_code}: {resp.text}"
+        )
+    try:
+        return json.loads(resp.text)
+    except json.JSONDecodeError:
+        LOGGER.error(f"Unable to parse {resp.text}")
+        raise
+
+
+def execute_model_registry_post_command(url: str, headers: dict[str, str], data_json: dict[str, Any]) -> dict[Any, Any]:
+    resp = requests.post(url=url, json=data_json, headers=headers, verify=False, timeout=60)
+    LOGGER.info(f"url: {url}, status code: {resp.status_code}, rep: {resp.text}")
+
+    if resp.status_code not in [200, 201]:
+        raise ModelRegistryResourceNotCreated(
+            f"Failed to create ModelRegistry resource: {url}, {resp.status_code}: {resp.text}"
+        )
+    try:
+        return json.loads(resp.text)
+    except json.JSONDecodeError:
+        LOGGER.error(f"Unable to parse {resp.text}")
+        raise
+
+
+def execute_model_registry_get_command(url: str, headers: dict[str, str]) -> dict[Any, Any]:  # skip-unused-code
+    resp = requests.get(url=url, headers=headers, verify=False)
+    LOGGER.info(f"url: {url}, status code: {resp.status_code}, rep: {resp.text}")
+    if resp.status_code not in [200, 201]:
+        raise ModelRegistryResourceNotFoundError(
+            f"Failed to get ModelRegistry resource: {url}, {resp.status_code}: {resp.text}"
+        )
+
+    try:
+        return json.loads(resp.text)
+    except json.JSONDecodeError:
+        LOGGER.error(f"Unable to parse {resp.text}")
+        raise
+
+
+def register_model_rest_api(
+    model_registry_rest_url: str, model_registry_rest_headers: dict[str, str], data_dict: dict[str, Any]
+) -> dict[str, Any]:
+    # register a model
+    register_model = execute_model_registry_post_command(
+        url=f"{model_registry_rest_url}{MODEL_REGISTRY_BASE_URI}registered_models",
+        headers=model_registry_rest_headers,
+        data_json=data_dict["register_model_data"],
+    )
+    # create associated model version:
+    model_data = data_dict["model_version_data"]
+    model_data["registeredModelId"] = register_model["id"]
+    model_version = execute_model_registry_post_command(
+        url=f"{model_registry_rest_url}{MODEL_REGISTRY_BASE_URI}model_versions",
+        headers=model_registry_rest_headers,
+        data_json=model_data,
+    )
+    # create associated model artifact
+    model_artifact = execute_model_registry_post_command(
+        url=f"{model_registry_rest_url}{MODEL_REGISTRY_BASE_URI}model_versions/{model_version['id']}/artifacts",
+        headers=model_registry_rest_headers,
+        data_json=data_dict["model_artifact_data"],
+    )
+    LOGGER.info(
+        f"Successfully registered model: {register_model}, with version: {model_version} and "
+        f"associated artifact: {model_artifact}"
+    )
+    return {"register_model": register_model, "model_version": model_version, "model_artifact": model_artifact}