Remove RHEC tests and add multiple source tests (#597)

dbasunag · web-flow · commit 3bb31da41768 · 2025-09-09T11:17:25.000-04:00
diff --git a/tests/model_registry/model_catalog/conftest.py b/tests/model_registry/model_catalog/conftest.py
@@ -8,7 +8,8 @@
 
 from ocp_resources.route import Route
 from tests.model_registry.constants import DEFAULT_MODEL_CATALOG
-from tests.model_registry.model_catalog.utils import is_model_catalog_ready, wait_for_model_catalog_api
+from tests.model_registry.model_catalog.constants import SAMPLE_MODEL_NAME3, CUSTOM_CATALOG_ID1
+from tests.model_registry.model_catalog.utils import is_model_catalog_ready, wait_for_model_catalog_api, get_model_str
 
 
 @pytest.fixture(scope="class")
@@ -47,7 +48,8 @@ def updated_catalog_config_map(
 ) -> Generator[ConfigMap, None, None]:
     patches = {"data": {"sources.yaml": request.param["sources_yaml"]}}
     if "sample_yaml" in request.param:
-        patches["data"]["sample-catalog.yaml"] = request.param["sample_yaml"]
+        for key in request.param["sample_yaml"]:
+            patches["data"][key] = request.param["sample_yaml"][key]
 
     with ResourceEditor(patches={catalog_config_map: patches}):
         is_model_catalog_ready(client=admin_client, model_registry_namespace=model_registry_namespace)
@@ -59,3 +61,21 @@ def updated_catalog_config_map(
 @pytest.fixture(scope="class")
 def expected_catalog_values(request: pytest.FixtureRequest) -> dict[str, str]:
     return request.param
+
+
+@pytest.fixture(scope="function")
+def update_configmap_data_add_model(
+    request: pytest.FixtureRequest,
+    catalog_config_map: ConfigMap,
+    model_registry_namespace: str,
+    admin_client: DynamicClient,
+    model_catalog_rest_url: list[str],
+    model_registry_rest_headers: dict[str, str],
+) -> Generator[ConfigMap, None, None]:
+    patches = catalog_config_map.instance.to_dict()
+    patches["data"][f"{CUSTOM_CATALOG_ID1.replace('_', '-')}.yaml"] += get_model_str(model=SAMPLE_MODEL_NAME3)
+    with ResourceEditor(patches={catalog_config_map: patches}):
+        is_model_catalog_ready(client=admin_client, model_registry_namespace=model_registry_namespace)
+        wait_for_model_catalog_api(url=model_catalog_rest_url[0], headers=model_registry_rest_headers)
+        yield catalog_config_map
+    is_model_catalog_ready(client=admin_client, model_registry_namespace=model_registry_namespace)
diff --git a/tests/model_registry/model_catalog/constants.py b/tests/model_registry/model_catalog/constants.py
@@ -1,44 +1,15 @@
-CUSTOM_ECHO_CATALOG_ID: str = "sample_rhec_catalog"
-CUSTOM_MODEL_NAME = "rhelai1/modelcar-granite-7b-starter"
-CUSTOM_ECOSYSTEM_CATALOG: str = f"""catalogs:
-- name: Red Hat Ecosystem Catalog
-  id: {CUSTOM_ECHO_CATALOG_ID}
-  type: rhec
-  enabled: true
-  properties:
-    models:
-    - {CUSTOM_MODEL_NAME}
-"""
-EXPECTED_ECHO_CATALOG_VALUES: dict[str, str] = {
-    "id": CUSTOM_ECHO_CATALOG_ID,
-    "model_name": f"{CUSTOM_MODEL_NAME}:latest",
-}
-CUSTOM_CATALOG_ID: str = "sample_custom_catalog"
-CUSTOM_CATALOG_WITH_FILE = f"""catalogs:
-- name: Sample Catalog
-  id: {CUSTOM_CATALOG_ID}
-  type: yaml
-  enabled: true
-  properties:
-    yamlCatalogPath: sample-catalog.yaml
-"""
-SAMPLE_CATALOG_FILE_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
-SAMPLE_CATALOG_YAML = f"""source: Hugging Face
-models:
-- name: {SAMPLE_CATALOG_FILE_NAME}
-  description: test description.
-  readme: |-
-    # test read me information
-  provider: Mistral AI
-  logo: temp placeholder logo
-  license: apache-2.0
-  licenseLink: https://www.apache.org/licenses/LICENSE-2.0.txt
-  libraryName: transformers
-  artifacts:
-    - uri: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3/resolve/main/consolidated.safetensors
-"""
-EXPECTED_CUSTOM_CATALOG_VALUES: dict[str, str] = {"id": CUSTOM_CATALOG_ID, "model_name": SAMPLE_CATALOG_FILE_NAME}
+CUSTOM_CATALOG_ID1: str = "sample_custom_catalog1"
+CUSTOM_CATALOG_ID2: str = "sample_custom_catalog2"
+SAMPLE_MODEL_NAME1 = "mistralai/Mistral-7B-Instruct-v0.3"
+
+SAMPLE_MODEL_NAME2 = "mistralai/Devstral-Small-2505"
+EXPECTED_CUSTOM_CATALOG_VALUES: list[dict[str, str]] = [{"id": CUSTOM_CATALOG_ID1, "model_name": SAMPLE_MODEL_NAME1}]
+MULTIPLE_CUSTOM_CATALOG_VALUES: list[dict[str, str]] = [
+    {"id": CUSTOM_CATALOG_ID1, "model_name": SAMPLE_MODEL_NAME1},
+    {"id": CUSTOM_CATALOG_ID2, "model_name": SAMPLE_MODEL_NAME2},
+]
 DEFAULT_CATALOG_NAME: str = "Default Catalog"
 DEFAULT_CATALOG_ID: str = "default_catalog"
 CATALOG_TYPE: str = "yaml"
 DEFAULT_CATALOG_FILE: str = "/default/default-catalog.yaml"
+SAMPLE_MODEL_NAME3 = "mistralai/Ministral-8B-Instruct-2410"
diff --git a/tests/model_registry/model_catalog/test_custom_model_catalog.py b/tests/model_registry/model_catalog/test_custom_model_catalog.py
@@ -1,29 +1,49 @@
 from tests.model_registry.model_catalog.constants import (
-    CUSTOM_ECOSYSTEM_CATALOG,
-    CUSTOM_CATALOG_WITH_FILE,
-    SAMPLE_CATALOG_YAML,
     EXPECTED_CUSTOM_CATALOG_VALUES,
-    EXPECTED_ECHO_CATALOG_VALUES,
+    CUSTOM_CATALOG_ID1,
+    SAMPLE_MODEL_NAME1,
+    CUSTOM_CATALOG_ID2,
+    SAMPLE_MODEL_NAME2,
+    MULTIPLE_CUSTOM_CATALOG_VALUES,
+    SAMPLE_MODEL_NAME3,
 )
 from ocp_resources.config_map import ConfigMap
 import pytest
 from simple_logger.logger import get_logger
 from typing import Self
 
-from tests.model_registry.model_catalog.utils import execute_get_command
+from tests.model_registry.model_catalog.utils import (
+    execute_get_command,
+    get_catalog_str,
+    get_sample_yaml_str,
+    ResourceNotFoundError,
+)
 
 LOGGER = get_logger(name=__name__)
 
 
 @pytest.mark.parametrize(
     "updated_catalog_config_map, expected_catalog_values",
     [
-        pytest.param({"sources_yaml": CUSTOM_ECOSYSTEM_CATALOG}, EXPECTED_ECHO_CATALOG_VALUES, id="rhec_test_catalog"),
         pytest.param(
-            {"sources_yaml": CUSTOM_CATALOG_WITH_FILE, "sample_yaml": SAMPLE_CATALOG_YAML},
+            {
+                "sources_yaml": get_catalog_str(ids=[CUSTOM_CATALOG_ID1]),
+                "sample_yaml": {"sample-custom-catalog1.yaml": get_sample_yaml_str(models=[SAMPLE_MODEL_NAME1])},
+            },
             EXPECTED_CUSTOM_CATALOG_VALUES,
             id="file_test_catalog",
         ),
+        pytest.param(
+            {
+                "sources_yaml": get_catalog_str(ids=[CUSTOM_CATALOG_ID1, CUSTOM_CATALOG_ID2]),
+                "sample_yaml": {
+                    "sample-custom-catalog1.yaml": get_sample_yaml_str(models=[SAMPLE_MODEL_NAME1]),
+                    "sample-custom-catalog2.yaml": get_sample_yaml_str(models=[SAMPLE_MODEL_NAME2]),
+                },
+            },
+            MULTIPLE_CUSTOM_CATALOG_VALUES,
+            id="file_test_catalog",
+        ),
     ],
     indirect=True,
 )
@@ -32,7 +52,7 @@
     "updated_catalog_config_map",
 )
 class TestModelCatalogCustom:
-    def test_model_custom_catalog_sources(
+    def test_model_custom_catalog_list_sources(
         self: Self,
         updated_catalog_config_map: tuple[ConfigMap, str, str],
         model_catalog_rest_url: list[str],
@@ -42,14 +62,19 @@ def test_model_custom_catalog_sources(
         """
         Validate sources api for model catalog
         """
-        result = execute_get_command(
-            url=f"{model_catalog_rest_url[0]}sources",
+        url = f"{model_catalog_rest_url[0]}sources"
+        results = execute_get_command(
+            url=url,
             headers=model_registry_rest_headers,
         )["items"]
-        assert len(result) == 1
-        assert result[0]["id"] == expected_catalog_values["id"]
 
-    def test_model_custom_catalog_models(
+        LOGGER.info(f"Results: for uri {url}: {results}")
+        assert len(results) == len(expected_catalog_values)
+        ids_from_query = [result_entry["id"] for result_entry in results]
+        ids_expected = [expected_entry["id"] for expected_entry in expected_catalog_values]
+        assert sorted(ids_from_query) == sorted(ids_expected), f"Expected: {expected_catalog_values}. Actual: {results}"
+
+    def test_model_custom_catalog_get_models_by_source(
         self: Self,
         updated_catalog_config_map: tuple[ConfigMap, str, str],
         model_catalog_rest_url: list[str],
@@ -59,11 +84,15 @@ def test_model_custom_catalog_models(
         """
         Validate models api for model catalog associated with a specific source
         """
-        result = execute_get_command(
-            url=f"{model_catalog_rest_url[0]}models?source={expected_catalog_values['id']}",
-            headers=model_registry_rest_headers,
-        )["items"]
-        assert result, f"Expected custom models to be present. Actual: {result}"
+        for expected_entry in expected_catalog_values:
+            LOGGER.info(f"Expected entry: {expected_entry}")
+            url = f"{model_catalog_rest_url[0]}models?source={expected_entry['id']}"
+            result = execute_get_command(
+                url=url,
+                headers=model_registry_rest_headers,
+            )["items"]
+            LOGGER.info(f"URL: {url} Result: {result}")
+            assert result, f"Expected custom models to be present. Actual: {result}"
 
     def test_model_custom_catalog_get_model_by_name(
         self: Self,
@@ -75,14 +104,18 @@ def test_model_custom_catalog_get_model_by_name(
         """
         Get Model by name associated with a specific source
         """
-        model_name = expected_catalog_values["model_name"]
-        result = execute_get_command(
-            url=f"{model_catalog_rest_url[0]}sources/{expected_catalog_values['id']}/models/{model_name}",
-            headers=model_registry_rest_headers,
-        )
-        assert result["name"] == model_name
+        for expected_entry in expected_catalog_values:
+            LOGGER.info(f"Expected entry: {expected_entry}")
+            model_name = expected_entry["model_name"]
+            url = f"{model_catalog_rest_url[0]}sources/{expected_entry['id']}/models/{model_name}"
+            result = execute_get_command(
+                url=url,
+                headers=model_registry_rest_headers,
+            )
+            LOGGER.info(f"URL: {url} Result: {result}")
+            assert result["name"] == model_name
 
-    def test_model_custom_catalog_get_artifact(
+    def test_model_custom_catalog_get_model_artifact(
         self: Self,
         updated_catalog_config_map: tuple[ConfigMap, str, str],
         model_catalog_rest_url: list[str],
@@ -92,11 +125,54 @@ def test_model_custom_catalog_get_artifact(
         """
         Get Model artifacts for model associated with specific source
         """
-        model_name = expected_catalog_values["model_name"]
-        artifacts = execute_get_command(
-            url=f"{model_catalog_rest_url[0]}sources/{expected_catalog_values['id']}/models/{model_name}/artifacts",
+        for expected_entry in expected_catalog_values:
+            LOGGER.info(f"Expected entry: {expected_entry}")
+
+            model_name = expected_entry["model_name"]
+            url = f"{model_catalog_rest_url[0]}sources/{expected_entry['id']}/models/{model_name}/artifacts"
+
+            artifacts = execute_get_command(
+                url=url,
+                headers=model_registry_rest_headers,
+            )["items"]
+            LOGGER.info(f"URL: {url} Result: {artifacts}")
+
+            assert artifacts, f"No artifacts found for {model_name}"
+            assert artifacts[0]["uri"]
+
+    @pytest.mark.dependency(name="test_model_custom_catalog_add_model")
+    def test_model_custom_catalog_add_model(
+        self: Self,
+        model_catalog_rest_url: list[str],
+        model_registry_rest_headers: dict[str, str],
+        expected_catalog_values: dict[str, str],
+        update_configmap_data_add_model: dict[str, str],
+    ):
+        """
+        Add a model to a source and ensure it is added to the catalog
+        """
+        url = f"{model_catalog_rest_url[0]}sources/{CUSTOM_CATALOG_ID1}/models/{SAMPLE_MODEL_NAME3}"
+        result = execute_get_command(
+            url=url,
             headers=model_registry_rest_headers,
-        )["items"]
+        )
+        LOGGER.info(f"URL: {url} Result: {result}")
+        assert result["name"] == SAMPLE_MODEL_NAME3
 
-        assert artifacts, f"No artifacts found for {model_name}"
-        assert artifacts[0]["uri"]
+    @pytest.mark.dependency(depends=["test_model_custom_catalog_add_model"])
+    def test_model_custom_catalog_remove_model(
+        self: Self,
+        model_catalog_rest_url: list[str],
+        model_registry_rest_headers: dict[str, str],
+        expected_catalog_values: dict[str, str],
+    ):
+        """
+        Add a model to a source and ensure it is added to the catalog
+        """
+        url = f"{model_catalog_rest_url[0]}sources/{CUSTOM_CATALOG_ID1}/models/{SAMPLE_MODEL_NAME3}"
+        with pytest.raises(ResourceNotFoundError):
+            result = execute_get_command(
+                url=url,
+                headers=model_registry_rest_headers,
+            )
+            LOGGER.info(f"URL: {url} Result: {result}")
diff --git a/tests/model_registry/model_catalog/utils.py b/tests/model_registry/model_catalog/utils.py
@@ -7,7 +7,6 @@
 import requests
 from timeout_sampler import retry
 
-from class_generator.parsers.explain_parser import ResourceNotFoundError
 from ocp_resources.pod import Pod
 from tests.model_registry.model_catalog.constants import (
     DEFAULT_CATALOG_NAME,
@@ -20,6 +19,10 @@
 LOGGER = get_logger(name=__name__)
 
 
+class ResourceNotFoundError(Exception):
+    pass
+
+
 def _execute_get_call(url: str, headers: dict[str, str], verify: bool | str = False) -> requests.Response:
     resp = requests.get(url=url, headers=headers, verify=verify, timeout=60)
     if resp.status_code not in [200, 201]:
@@ -87,3 +90,47 @@ def validate_default_catalog(default_catalog) -> None:
     assert default_catalog["id"] == DEFAULT_CATALOG_ID
     assert default_catalog["type"] == CATALOG_TYPE
     assert default_catalog["properties"].get("yamlCatalogPath") == DEFAULT_CATALOG_FILE
+
+
+def get_catalog_str(ids: list[str]) -> str:
+    catalog_str: str = ""
+    for id in ids:
+        catalog_str += f"""
+- name: Sample Catalog
+  id: {id}
+  type: yaml
+  enabled: true
+  properties:
+    yamlCatalogPath: {id.replace("_", "-")}.yaml
+"""
+    return f"""catalogs:
+{catalog_str}
+    """
+
+
+def get_sample_yaml_str(models: list[str]) -> str:
+    model_str: str = ""
+    for model in models:
+        model_str += f"""
+{get_model_str(model=model)}
+"""
+    return f"""source: Hugging Face
+models:
+{model_str}
+"""
+
+
+def get_model_str(model: str) -> str:
+    return f"""
+- name: {model}
+  description: test description.
+  readme: |-
+    # test read me information {model}
+  provider: Mistral AI
+  logo: temp placeholder logo
+  license: apache-2.0
+  licenseLink: https://www.apache.org/licenses/LICENSE-2.0.txt
+  libraryName: transformers
+  artifacts:
+    - uri: https://huggingface.co/{model}/resolve/main/consolidated.safetensors
+"""
