Upgrade support for custom MC scenario (#736)

* Upgrade support for custom MC scenario

* updates! 2c9819d

* Add custom catalog tests

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

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

* Address review comments

---------

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

Author: dbasunag

conftest.py

@@ -247,15 +247,17 @@ def _add_upgrade_test(_item: Item, _upgrade_deployment_modes: list[str]) -> bool
             # Add support to be able to reuse tests in both upgrade and fresh install scenarios
             if "install" in item.keywords:
                 non_upgrade_tests.append(item)
-
+            if "post_upgrade" in item.keywords:
+                post_upgrade_tests.append(item)
         elif "post_upgrade" in item.keywords and _add_upgrade_test(
             _item=item, _upgrade_deployment_modes=upgrade_deployment_modes
         ):
             post_upgrade_tests.append(item)
             # Add support to be able to reuse tests in both upgrade and fresh install scenarios
             if "install" in item.keywords:
                 non_upgrade_tests.append(item)
-
+            if "pre_upgrade" in item.keywords:
+                pre_upgrade_tests.append(item)
         else:
             non_upgrade_tests.append(item)
 

tests/model_registry/conftest.py

@@ -13,6 +13,7 @@
 from ocp_resources.infrastructure import Infrastructure
 from ocp_resources.oauth import OAuth
 from ocp_resources.pod import Pod
+from ocp_resources.route import Route
 from ocp_resources.secret import Secret
 from ocp_resources.namespace import Namespace
 from ocp_resources.service import Service
@@ -43,6 +44,7 @@
     DB_RESOURCE_NAME,
     MR_INSTANCE_NAME,
     MODEL_REGISTRY_POD_FILTER,
+    DEFAULT_CUSTOM_MODEL_CATALOG,
     KUBERBACPROXY_STR,
 )
 from utilities.constants import Labels, Protocols
@@ -591,3 +593,28 @@ def user_credentials_rbac() -> dict[str, str]:
         "idp_name": f"test-htpasswd-idp-{random_str}",
         "secret_name": f"test-htpasswd-secret-{random_str}",
     }
+
+
+@pytest.fixture(scope="class")
+def catalog_config_map(admin_client: DynamicClient, model_registry_namespace: str) -> ConfigMap:
+    return ConfigMap(name=DEFAULT_CUSTOM_MODEL_CATALOG, client=admin_client, namespace=model_registry_namespace)
+
+
+@pytest.fixture(scope="class")
+def model_catalog_routes(admin_client: DynamicClient, model_registry_namespace: str) -> list[Route]:
+    return list(
+        Route.get(namespace=model_registry_namespace, label_selector="component=model-catalog", dyn_client=admin_client)
+    )
+
+
+@pytest.fixture(scope="class")
+def model_catalog_rest_url(model_registry_namespace: str, model_catalog_routes: list[Route]) -> list[str]:
+    assert model_catalog_routes, f"Model catalog routes does not exist in {model_registry_namespace}"
+    route_urls = [
+        f"https://{route.instance.spec.host}:443/api/model_catalog/v1alpha1/" for route in model_catalog_routes
+    ]
+    assert route_urls, (
+        "Model catalog routes information could not be found from "
+        f"routes:{[route.name for route in model_catalog_routes]}"
+    )
+    return route_urls

tests/model_registry/constants.py

@@ -65,5 +65,7 @@ class ModelRegistryEndpoints:
 }
 MODEL_REGISTRY_POD_FILTER: str = "component=model-registry"
 DEFAULT_CUSTOM_MODEL_CATALOG: str = "model-catalog-sources"
+SAMPLE_MODEL_NAME1 = "mistralai/Mistral-7B-Instruct-v0.3"
+CUSTOM_CATALOG_ID1: str = "sample_custom_catalog1"
 DEFAULT_MODEL_CATALOG_CM: str = "model-catalog-default-sources"
 KUBERBACPROXY_STR: str = "KubeRBACProxyAvailable"

tests/model_registry/model_catalog/conftest.py

@@ -10,36 +10,29 @@
 from ocp_resources.config_map import ConfigMap
 from ocp_resources.resource import ResourceEditor
 
-from ocp_resources.route import Route
 from ocp_resources.service_account import ServiceAccount
-from tests.model_registry.constants import DEFAULT_CUSTOM_MODEL_CATALOG
 from tests.model_registry.model_catalog.constants import (
     SAMPLE_MODEL_NAME3,
-    CUSTOM_CATALOG_ID1,
     DEFAULT_CATALOG_FILE,
     CATALOG_CONTAINER,
     REDHAT_AI_CATALOG_ID,
 )
-from tests.model_registry.model_catalog.utils import (
+from tests.model_registry.constants import CUSTOM_CATALOG_ID1
+from tests.model_registry.utils import (
+    get_rest_headers,
     is_model_catalog_ready,
+    get_model_catalog_pod,
     wait_for_model_catalog_api,
-    get_model_str,
     execute_get_command,
-    get_model_catalog_pod,
+    get_model_str,
 )
-from tests.model_registry.utils import get_rest_headers
 from utilities.infra import get_openshift_token, login_with_user_password, create_inference_token
 from utilities.user_utils import UserTestSession
 
 
 LOGGER = get_logger(name=__name__)
 
 
-@pytest.fixture(scope="class")
-def catalog_config_map(admin_client: DynamicClient, model_registry_namespace: str) -> ConfigMap:
-    return ConfigMap(name=DEFAULT_CUSTOM_MODEL_CATALOG, client=admin_client, namespace=model_registry_namespace)
-
-
 @pytest.fixture(scope="class")
 def model_catalog_config_map(
     request: pytest.FixtureRequest, admin_client: DynamicClient, model_registry_namespace: str
@@ -50,26 +43,6 @@ def model_catalog_config_map(
     return ConfigMap(name=configmap_name, client=admin_client, namespace=model_registry_namespace, ensure_exists=True)
 
 
-@pytest.fixture(scope="class")
-def model_catalog_routes(admin_client: DynamicClient, model_registry_namespace: str) -> list[Route]:
-    return list(
-        Route.get(namespace=model_registry_namespace, label_selector="component=model-catalog", dyn_client=admin_client)
-    )
-
-
-@pytest.fixture(scope="class")
-def model_catalog_rest_url(model_registry_namespace: str, model_catalog_routes: list[Route]) -> list[str]:
-    assert model_catalog_routes, f"Model catalog routes does not exist in {model_registry_namespace}"
-    route_urls = [
-        f"https://{route.instance.spec.host}:443/api/model_catalog/v1alpha1/" for route in model_catalog_routes
-    ]
-    assert route_urls, (
-        "Model catalog routes information could not be found from "
-        f"routes:{[route.name for route in model_catalog_routes]}"
-    )
-    return route_urls
-
-
 @pytest.fixture(scope="class")
 def updated_catalog_config_map(
     request: pytest.FixtureRequest,

tests/model_registry/model_catalog/constants.py

@@ -1,8 +1,8 @@
 from typing import Any
 
-CUSTOM_CATALOG_ID1: str = "sample_custom_catalog1"
+from tests.model_registry.constants import SAMPLE_MODEL_NAME1, CUSTOM_CATALOG_ID1
+
 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}]

tests/model_registry/model_catalog/test_custom_model_catalog.py

@@ -1,22 +1,22 @@
 from tests.model_registry.model_catalog.constants import (
     EXPECTED_CUSTOM_CATALOG_VALUES,
-    CUSTOM_CATALOG_ID1,
-    SAMPLE_MODEL_NAME1,
     CUSTOM_CATALOG_ID2,
     SAMPLE_MODEL_NAME2,
     MULTIPLE_CUSTOM_CATALOG_VALUES,
     SAMPLE_MODEL_NAME3,
 )
+from tests.model_registry.constants import SAMPLE_MODEL_NAME1, CUSTOM_CATALOG_ID1
 from ocp_resources.config_map import ConfigMap
 import pytest
 from simple_logger.logger import get_logger
 from typing import Self
+from kubernetes.dynamic.exceptions import ResourceNotFoundError
 
-from tests.model_registry.model_catalog.utils import (
+from tests.model_registry.utils import (
     execute_get_command,
-    get_catalog_str,
     get_sample_yaml_str,
-    ResourceNotFoundError,
+    get_catalog_str,
+    validate_model_catalog_sources,
 )
 
 LOGGER = get_logger(name=__name__)
@@ -32,6 +32,7 @@
             },
             EXPECTED_CUSTOM_CATALOG_VALUES,
             id="test_file_test_catalog",
+            marks=(pytest.mark.pre_upgrade, pytest.mark.post_upgrade, pytest.mark.install),
         ),
         pytest.param(
             {
@@ -62,15 +63,10 @@ def test_model_custom_catalog_list_sources(
         """
         Validate sources api for model catalog
         """
-        url = f"{model_catalog_rest_url[0]}sources"
-        results = execute_get_command(
-            url=url,
-            headers=model_registry_rest_headers,
-        )["items"]
-        ids_from_query = [result_entry["id"] for result_entry in results]
-        ids_expected = [expected_entry["id"] for expected_entry in expected_catalog_values]
-        assert set(ids_expected).issubset(set(ids_from_query)), (
-            f"Expected model catalogs: {expected_catalog_values}. Actual model catalogs: {results}"
+        validate_model_catalog_sources(
+            model_catalog_sources_url=f"{model_catalog_rest_url[0]}sources",
+            rest_headers=model_registry_rest_headers,
+            expected_catalog_values=expected_catalog_values,
         )
 
     def test_model_custom_catalog_get_models_by_source(

tests/model_registry/model_catalog/test_default_model_catalog.py

@@ -1,5 +1,7 @@
 import pytest
 import random
+
+import yaml
 from kubernetes.dynamic import DynamicClient
 from dictdiffer import diff
 from ocp_resources.deployment import Deployment
@@ -10,18 +12,17 @@
 from ocp_resources.config_map import ConfigMap
 from ocp_resources.route import Route
 from ocp_resources.service import Service
+
 from tests.model_registry.constants import DEFAULT_CUSTOM_MODEL_CATALOG, DEFAULT_MODEL_CATALOG_CM
-from tests.model_registry.model_catalog.constants import REDHAT_AI_CATALOG_ID, CATALOG_CONTAINER
+from tests.model_registry.model_catalog.constants import REDHAT_AI_CATALOG_ID, CATALOG_CONTAINER, DEFAULT_CATALOGS
 from tests.model_registry.model_catalog.utils import (
     validate_model_catalog_enabled,
-    execute_get_command,
     validate_model_catalog_resource,
     get_validate_default_model_catalog_source,
     extract_schema_fields,
-    get_model_catalog_pod,
-    validate_model_catalog_configmap_data,
+    validate_default_catalog,
 )
-from tests.model_registry.utils import get_rest_headers
+from tests.model_registry.utils import get_rest_headers, get_model_catalog_pod, execute_get_command
 from utilities.user_utils import UserTestSession
 
 LOGGER = get_logger(name=__name__)
@@ -36,23 +37,41 @@
 @pytest.mark.skip_must_gather
 class TestModelCatalogGeneral:
     @pytest.mark.parametrize(
-        "model_catalog_config_map, expected_catalogs",
+        "model_catalog_config_map, expected_catalogs, validate_catalog",
         [
             pytest.param(
                 {"configmap_name": DEFAULT_CUSTOM_MODEL_CATALOG},
                 0,
-                id="test_model_catalog_sources_configmap",
+                False,
+                id="test_model_catalog_sources_configmap_install",
+                marks=pytest.mark.install,
+            ),
+            pytest.param(
+                {"configmap_name": DEFAULT_CUSTOM_MODEL_CATALOG},
+                1,
+                False,
+                id="test_model_catalog_sources_configmap_upgrade",
+                marks=(pytest.mark.pre_upgrade, pytest.mark.post_upgrade),
             ),
             pytest.param(
                 {"configmap_name": DEFAULT_MODEL_CATALOG_CM},
                 2,
+                True,
                 id="test_model_catalog_default_sources_configmap",
             ),
         ],
         indirect=["model_catalog_config_map"],
     )
-    def test_config_map_exists(self: Self, model_catalog_config_map: ConfigMap, expected_catalogs: int):
-        validate_model_catalog_configmap_data(configmap=model_catalog_config_map, num_catalogs=expected_catalogs)
+    def test_config_map_exists(
+        self: Self, model_catalog_config_map: ConfigMap, expected_catalogs: int, validate_catalog: bool
+    ) -> None:
+        assert model_catalog_config_map.exists, f"{model_catalog_config_map.name} does not exist"
+        catalogs = yaml.safe_load(model_catalog_config_map.instance.data["sources.yaml"])["catalogs"]
+        assert len(catalogs) == expected_catalogs, (
+            f"{model_catalog_config_map.name} should have {expected_catalogs} catalog"
+        )
+        if validate_catalog:
+            validate_default_catalog(catalogs=catalogs)
 
     @pytest.mark.parametrize(
         "resource_name, expected_resource_count",
@@ -80,6 +99,7 @@ def test_config_map_exists(self: Self, model_catalog_config_map: ConfigMap, expe
         ],
     )
     @pytest.mark.post_upgrade
+    @pytest.mark.pre_upgrade
     @pytest.mark.install
     def test_model_catalog_resources_exists(
         self: Self,
@@ -95,9 +115,15 @@ def test_model_catalog_resources_exists(
             expected_resource_count=expected_resource_count,
         )
 
+    @pytest.mark.post_upgrade
+    @pytest.mark.pre_upgrade
+    @pytest.mark.install
     def test_operator_pod_enabled_model_catalog(self: Self, model_registry_operator_pod: Pod):
         assert validate_model_catalog_enabled(pod=model_registry_operator_pod)
 
+    @pytest.mark.post_upgrade
+    @pytest.mark.pre_upgrade
+    @pytest.mark.install
     def test_model_catalog_uses_postgres(self: Self, admin_client: DynamicClient, model_registry_namespace: str):
         """
         Validate that model catalog pod is using PostgreSQL database
@@ -119,6 +145,7 @@ def test_model_catalog_uses_postgres(self: Self, admin_client: DynamicClient, mo
         pytest.param(
             {},
             id="test_model_catalog_source_admin_user",
+            marks=(pytest.mark.pre_upgrade, pytest.mark.post_upgrade, pytest.mark.install),
         ),
         pytest.param(
             {"user_type": "test"},
@@ -134,16 +161,29 @@ def test_model_catalog_uses_postgres(self: Self, admin_client: DynamicClient, mo
 class TestModelCatalogDefault:
     def test_model_catalog_default_catalog_sources(
         self,
+        pytestconfig: pytest.Config,
         test_idp_user: UserTestSession,
         model_catalog_rest_url: list[str],
         user_token_for_api_calls: str,
     ):
         """
         Validate specific user can access default model catalog source
         """
-        get_validate_default_model_catalog_source(
-            token=user_token_for_api_calls, model_catalog_url=f"{model_catalog_rest_url[0]}sources"
-        )
+        LOGGER.info("Attempting client connection with token")
+        result = execute_get_command(
+            url=f"{model_catalog_rest_url[0]}sources",
+            headers=get_rest_headers(token=user_token_for_api_calls),
+        )["items"]
+        assert result
+        items_to_validate = []
+        if pytestconfig.option.pre_upgrade or pytestconfig.option.post_upgrade:
+            for catalog in result:
+                if catalog["id"] in DEFAULT_CATALOGS.keys():
+                    items_to_validate.append(catalog)
+            assert len(items_to_validate) + 1 == len(result)
+        else:
+            items_to_validate = result
+        get_validate_default_model_catalog_source(catalogs=items_to_validate)
 
     def test_model_default_catalog_get_models_by_source(
         self: Self,
@@ -191,6 +231,9 @@ def test_model_default_catalog_get_model_artifact(
         assert result[0]["uri"]
 
 
+@pytest.mark.post_upgrade
+@pytest.mark.pre_upgrade
+@pytest.mark.install
 @pytest.mark.skip_must_gather
 class TestModelCatalogDefaultData:
     """Test class for validating default catalog data (not user-specific)"""

tests/model_registry/model_catalog/test_filter_options_endpoint.py

@@ -2,8 +2,8 @@
 from typing import Self
 from simple_logger.logger import get_logger
 
-from tests.model_registry.model_catalog.utils import execute_get_command, validate_filter_options_structure
-from tests.model_registry.utils import get_rest_headers
+from tests.model_registry.model_catalog.utils import validate_filter_options_structure
+from tests.model_registry.utils import get_rest_headers, execute_get_command
 from utilities.user_utils import UserTestSession
 
 LOGGER = get_logger(name=__name__)

tests/model_registry/model_catalog/test_model_catalog_negative.py

@@ -9,8 +9,8 @@
 from ocp_resources.resource import ResourceEditor
 from tests.model_registry.constants import DEFAULT_MODEL_CATALOG_CM
 from tests.model_registry.model_catalog.constants import DEFAULT_CATALOGS, CATALOG_CONTAINER
-from tests.model_registry.model_catalog.utils import validate_model_catalog_configmap_data, is_model_catalog_ready
-from tests.model_registry.utils import get_model_catalog_pod
+from tests.model_registry.model_catalog.utils import validate_model_catalog_configmap_data
+from tests.model_registry.utils import get_model_catalog_pod, is_model_catalog_ready
 from timeout_sampler import TimeoutExpiredError
 
 from utilities.general import wait_for_container_status

tests/model_registry/model_catalog/test_model_search.py

@@ -9,7 +9,7 @@
     REDHAT_AI_CATALOG_ID,
     VALIDATED_CATALOG_ID,
 )
-from tests.model_registry.model_catalog.utils import (
+from tests.model_registry.utils import (
     execute_get_command,
 )