Add tests for model_artifact update validations

[dbasunag]

Description

How Has This Been Tested?

Merge criteria:

• The commits are squashed in a cohesive manner and have meaningful messages.
• Testing instructions have been added in the PR body (for PRs involving changes that are not immediately obvious).
• The developer has manually tested the changes and verified that the changes work

Summary by CodeRabbit

• New Features

  • Introduced comprehensive tests for model registry REST API creation and update operations.
  • Added fixtures and utilities for streamlined model registration and artifact update testing.
  • Provided structured test data and constants for model registry resources.
  • Added custom exceptions for precise error handling in model registry operations.

• Bug Fixes

  • Improved error handling and reporting for REST API interactions within tests.

• Chores

  • Enhanced test coverage and organization for model registry REST API functionality.

[coderabbitai]

"""

Walkthrough

This update introduces a comprehensive suite of test infrastructure for a model registry REST API. It adds custom exception classes, constants for test data, REST API utility functions, pytest fixtures, and a test class with methods to validate model registry creation and update operations. The new modules provide structured data, error handling, and reusable fixtures to facilitate robust testing of model registration, versioning, and artifact management through HTTP requests. The changes collectively establish a foundation for automated testing of model registry REST endpoints.

Changes

File(s)	Change Summary
tests/model_registry/exceptions.py	Added three custom exception classes: ModelRegistryResourceNotCreated, ModelRegistryResourceNotFoundError, and ModelRegistryResourceNotUpdated, each inheriting from Exception and serving as specific error types for model registry operations.
tests/model_registry/rest_api/constants.py	Introduced constants for test data: model registration, version, artifact dictionaries, an aggregate data dictionary, and a base URI string. Imports model format enumeration for artifact format specification.
tests/model_registry/rest_api/utils.py	Added utility functions for executing PATCH, POST, and GET HTTP requests with error handling and logging. Introduced register_model_rest_api to orchestrate model, version, and artifact registration via REST API calls, returning structured response data.
tests/model_registry/rest_api/conftest.py	Added pytest fixtures for constructing REST URLs, preparing HTTP headers, registering models via REST API, and updating model artifacts. Fixtures use utility functions and constants, and provide reusable setup for REST API tests.
tests/model_registry/rest_api/test_model_registry_rest_api.py	Added test class TestModelRegistryCreationRest with methods to validate model registry resource creation and artifact update operations. Tests use parameterization and fixtures to compare expected and actual resource attributes, with structured error reporting.

Sequence Diagram(s)

sequenceDiagram
    participant Test as Test Method
    participant Fixture as Pytest Fixture
    participant Utils as REST API Utils
    participant API as Model Registry REST API

    Test->>Fixture: Request registered_model_rest_api
    Fixture->>Utils: register_model_rest_api(...)
    Utils->>API: POST /registered_model
    API-->>Utils: Registered Model Response
    Utils->>API: POST /model_version (with registered_model_id)
    API-->>Utils: Model Version Response
    Utils->>API: POST /model_artifact (with model_version_id)
    API-->>Utils: Model Artifact Response
    Utils-->>Fixture: Combined Response
    Fixture-->>Test: Registered Model, Version, Artifact

    Test->>Fixture: Request updated_model_artifact
    Fixture->>Utils: execute_model_registry_patch_command(...)
    Utils->>API: PATCH /model_artifact/{id}
    API-->>Utils: Updated Artifact Response
    Utils-->>Fixture: Updated Artifact
    Fixture-->>Test: Updated Artifact

Loading

Poem

In the warren where code bunnies hop,
We built REST tests that simply won’t stop!
Exceptions are custom, our fixtures robust,
Constants and utils—precision’s a must.
Now models and versions leap into place,
Artifacts updated, all passing the race!
🐇✨
"""

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 320bfc6 and 39cbba1.

📒 Files selected for processing (5)

• tests/model_registry/exceptions.py (1 hunks)
• tests/model_registry/rest_api/conftest.py (1 hunks)
• tests/model_registry/rest_api/constants.py (1 hunks)
• tests/model_registry/rest_api/test_model_registry_rest_api.py (1 hunks)
• tests/model_registry/rest_api/utils.py (1 hunks)

✅ Files skipped from review due to trivial changes (2)

• tests/model_registry/rest_api/constants.py
• tests/model_registry/rest_api/conftest.py

🚧 Files skipped from review as they are similar to previous changes (3)

• tests/model_registry/exceptions.py
• tests/model_registry/rest_api/test_model_registry_rest_api.py
• tests/model_registry/rest_api/utils.py

✨ Finishing Touches

• 📝 Generate Docstrings

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[github-actions]

The following are automatically added/executed:

• PR size label.
• Run pre-commit
• Run tox
• Add PR author as the PR assignee

Available user actions:

• To mark a PR as WIP, add /wip in a comment. To remove it from the PR comment /wip cancel to the PR.
• To block merging of a PR, add /hold in a comment. To un-block merging of PR comment /hold cancel.
• To mark a PR as approved, add /lgtm in a comment. To remove, add /lgtm cancel.
  lgtm label removed on each new commit push.
• To mark PR as verified comment /verified to the PR, to un-verify comment /verified cancel to the PR.
  verified label removed on each new commit push.
• To Cherry-pick a merged PR /cherry-pick <target_branch_name> to the PR. If <target_branch_name> is valid,
  and the current PR is merged, a cherry-picked PR would be created and linked to the current PR.

Supported labels

{'/verified', '/hold', '/wip', '/lgtm'}

[coderabbitai]

Actionable comments posted: 6

🧹 Nitpick comments (10)

tests/model_registry/exceptions.py (1)

11-11: Add a trailing newline at the end of the file.

The file is missing a trailing newline at the end, which is recommended by PEP 8 and prevents warnings from many linters.

 class ModelRegistryResourceNotUpdated(Exception):
     pass
-
+

tests/model_registry/rest_api/constants.py (1)

31-32: Add a trailing newline at the end of the file.

The file is missing a trailing newline at the end, which is recommended by PEP 8 and prevents warnings from many linters.

 MODEL_REGISTRY_BASE_URI = "/api/model_registry/v1alpha3/"
-
+

tests/model_registry/rest_api/test_model_registry_rest_api.py (5)

61-61: Fix the type annotation for errors list.

The type annotation for the errors list is incorrect. It should be list[dict[str, Any]] since you're creating a list of dictionaries in the comprehension.

-        errors: list[str:Any]
+        errors: list[dict[str, Any]]

---

97-97: Fix the type annotation for errors list in the update test.

The type annotation for the errors list should include the nested type structure.

-        errors: list[dict[str, list[Any]]]
+        errors: list[dict[str, list[Any]]]

---

55-69: Consider extracting the validation logic to a helper method.

Both test methods use similar validation logic with list comprehensions to check if expected values match actual values. Consider extracting this to a helper method to avoid duplication.

+    def _validate_expected_values(
+        self,
+        actual_data: dict[str, Any],
+        expected_params: dict[str, Any],
+        resource_name: str
+    ) -> None:
+        errors: list[dict[str, list[Any]]] = [
+            {key: [expected_params[key], actual_data.get(key)]}
+            for key in expected_params.keys()
+            if not actual_data.get(key) or actual_data[key] != expected_params[key]
+        ]
+        if errors:
+            pytest.fail(f"{resource_name} did not get created/updated with expected values: {errors}")
+        LOGGER.info(f"Successfully validated: {actual_data['name']}")
+
     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], registered_model_rest_api.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']}")
+        self._validate_expected_values(
+            actual_data=created_resource_data, 
+            expected_params=expected_params,
+            resource_name="Model"
+        )

---

92-104: Reuse validation logic in second test method.

The same validation logic from the first test method is duplicated here. Update this method to use the suggested helper method.

    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']}")
+        self._validate_expected_values(
+            actual_data=updated_model_artifact,
+            expected_params=expected_param,
+            resource_name="Model artifact"
+        )

---

105-105: Add a trailing newline at the end of the file.

The file is missing a trailing newline at the end, which is recommended by PEP 8 and prevents warnings from many linters.

        LOGGER.info(f"Successfully validated: {updated_model_artifact['name']}")
-
+

tests/model_registry/rest_api/conftest.py (1)

50-50: Add a trailing newline at the end of the file.

The file is missing a trailing newline at the end, which is recommended by PEP 8 and prevents warnings from many linters.

    )
-
+

tests/model_registry/rest_api/utils.py (2)

62-89: Add input validation and optimize return structure.

The function assumes the correct structure of data_dict without validating it. Consider adding validation for required keys before using them. Also, the return structure could be improved for more consistent access.

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]:
+    # Validate required keys in data_dict
+    required_keys = ["register_model_data", "model_version_data", "model_artifact_data"]
+    for key in required_keys:
+        if key not in data_dict:
+            raise ValueError(f"Missing required key in data_dict: {key}")
+
    # 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}
+    return {
+        "register_model": register_model,
+        "model_version": model_version,
+        "model_artifact": model_artifact
+    }

---

1-90: Consider refactoring HTTP request functions to reduce duplication.

There's significant code duplication among the three HTTP request functions. Consider refactoring to extract common logic into a shared helper function.

def _execute_model_registry_request(
    method: str, url: str, headers: dict[str, str], data_json: dict[str, Any] = None, 
    expected_status_codes: list[int] = None
) -> dict[Any, Any]:
    """
    Helper function to execute HTTP requests against the model registry API.
    
    Args:
        method: HTTP method (get, post, patch)
        url: Request URL
        headers: Request headers
        data_json: JSON data for POST/PATCH requests
        expected_status_codes: Status codes considered as successful
        
    Returns:
        Parsed JSON response
        
    Raises:
        ModelRegistryResourceNotCreated: For POST failures
        ModelRegistryResourceNotUpdated: For PATCH failures
        ModelRegistryResourceNotFoundError: For GET failures
    """
    if expected_status_codes is None:
        expected_status_codes = [200]
        
    request_fn = getattr(requests, method.lower())
    kwargs = {"url": url, "headers": headers, "verify": False, "timeout": 30}
    if method.lower() in ["post", "patch"] and data_json is not None:
        kwargs["json"] = data_json
        
    resp = request_fn(**kwargs)
    LOGGER.info(f"url: {url}, status code: {resp.status_code}, rep: {resp.text}")
    
    if resp.status_code not in expected_status_codes:
        if method.lower() == "post":
            raise ModelRegistryResourceNotCreated(
                f"Failed to create ModelRegistry resource: {url}, {resp.status_code}: {resp.text}"
            )
        elif method.lower() == "patch":
            raise ModelRegistryResourceNotUpdated(
                f"Failed to update ModelRegistry resource: {url}, {resp.status_code}: {resp.text}"
            )
        else:  # get
            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

Then the individual functions would become much simpler:

def execute_model_registry_patch_command(
    url: str, headers: dict[str, str], data_json: dict[str, Any]
) -> dict[Any, Any]:
    return _execute_model_registry_request("patch", url, headers, data_json)

def execute_model_registry_post_command(
    url: str, headers: dict[str, str], data_json: dict[str, Any]
) -> dict[Any, Any]:
    return _execute_model_registry_request("post", url, headers, data_json, [200, 201])

def execute_model_registry_get_command(url: str, headers: dict[str, str]) -> dict[Any, Any]:
    return _execute_model_registry_request("get", url, headers)

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 560f952 and 2a936ca.

📒 Files selected for processing (5)

• tests/model_registry/exceptions.py (1 hunks)
• tests/model_registry/rest_api/conftest.py (1 hunks)
• tests/model_registry/rest_api/constants.py (1 hunks)
• tests/model_registry/rest_api/test_model_registry_rest_api.py (1 hunks)
• tests/model_registry/rest_api/utils.py (1 hunks)

🧰 Additional context used

🧬 Code Graph Analysis (4)

tests/model_registry/rest_api/constants.py (1)

utilities/constants.py (1)

• ModelFormat (12-19)

tests/model_registry/rest_api/conftest.py (4)

tests/model_registry/rest_api/utils.py (2)

• register_model_rest_api (62-89)
• execute_model_registry_patch_command (15-29)

utilities/constants.py (1)

• Protocols (87-94)

tests/model_registry/conftest.py (1)

• model_registry_instance_rest_endpoint (188-191)

tests/conftest.py (1)

• current_client_token (81-82)

tests/model_registry/rest_api/test_model_registry_rest_api.py (2)

utilities/constants.py (2)

• DscComponents (146-164)
• ManagementState (151-153)

tests/model_registry/rest_api/conftest.py (2)

• registered_model_rest_api (26-33)
• updated_model_artifact (37-49)

tests/model_registry/rest_api/utils.py (2)

tests/model_registry/exceptions.py (3)

• ModelRegistryResourceNotCreated (1-2)
• ModelRegistryResourceNotFoundError (5-6)
• ModelRegistryResourceNotUpdated (9-10)

tests/model_registry/rest_api/conftest.py (2)

• model_registry_rest_url (11-13)
• model_registry_rest_headers (17-22)

🔇 Additional comments (13)

tests/model_registry/exceptions.py (1)

1-10: Clean exception hierarchy for model registry operations.

These custom exception classes provide clear error signaling for different failure scenarios in the model registry operations. They will help distinguish between creation, lookup, and update failures during testing.

tests/model_registry/rest_api/constants.py (4)

5-9: LGTM: Clear and descriptive model registration data.

The model registration data dictionary has appropriate fields with descriptive values.

---

10-15: LGTM: Well-structured model version data.

The model version data is properly structured with appropriate state and metadata fields.

---

17-25: LGTM: Comprehensive model artifact test data.

The model artifact data contains all necessary fields including URI, format details, and state. Using the ModelFormat enum from utilities.constants ensures consistency with the rest of the codebase.

---

26-31: LGTM: Convenient aggregation of test data.

Aggregating the different data dictionaries into MODEL_REGISTER_DATA provides a convenient way to pass all the data to the test fixtures at once.

tests/model_registry/rest_api/test_model_registry_rest_api.py (4)

11-28: LGTM: Well-structured test parametrization for model registry configuration.

The parametrization correctly sets up the test environment with the model registry component configured as "Managed" with a specific namespace.

---

29-33: LGTM: Clear test class docstring.

The docstring clearly explains the purpose of the test class and its behavior regarding component state management.

---

35-54: LGTM: Comprehensive test parametrization for model registry resources.

The parametrization covers all three resource types (model register, version, and artifact) with expected parameter values for validation.

---

71-91: LGTM: Well-structured test parameters for model artifact updates.

The parametrization covers multiple update scenarios for the model artifact with different fields (description, format name, format version).

tests/model_registry/rest_api/conftest.py (4)

10-14: LGTM: Well-structured fixture for model registry REST URL.

The fixture correctly builds the REST URL using the HTTPS protocol and the model registry endpoint.

---

16-23: LGTM: Comprehensive REST headers fixture.

The headers fixture includes all necessary fields for API authentication and content negotiation.

---

25-34: LGTM: Clear model registration fixture.

The fixture correctly utilizes the utility function to register a model via the REST API and passes through the test parameters.

---

36-49: LGTM: Well-implemented model artifact update fixture.

The fixture properly extracts the model artifact ID, validates its existence, and executes the PATCH command to update the artifact. The assertion on line 44 provides good error handling.

[lugi0]

Why did you decide to use the API directly rather than the Python client?

[dbasunag]

Why did you decide to use the API directly rather than the Python client?

Python client has a few restrictions.
e.g.it would create a resource if it does not exist and update an existing resource otherwise. In testing, if the resource exist we should fail. But using python client we would never know if it is updating a resource vs creating a new resource.

[dbasunag]

/verified

[fege]

/lgtm

[github-actions]

Status of building tag latest: success.
Status of pushing tag latest to image registry: failure.