Add custom catalog upgrade scenarios

Author: dbasunag

conftest.py

@@ -247,14 +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_MODEL_CATALOG,
 )
 from utilities.constants import Labels, Protocols
 from tests.model_registry.utils import (
@@ -590,3 +592,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 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 catalog_config_map(admin_client: DynamicClient, model_registry_namespace: str) -> ConfigMap:
+    return ConfigMap(name=DEFAULT_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)
+    )

tests/model_registry/constants.py

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

tests/model_registry/model_catalog/conftest.py

@@ -1,5 +1,6 @@
 import random
 from typing import Generator, Any
+
 from simple_logger.logger import get_logger
 
 import pytest
@@ -8,67 +9,47 @@
 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_MODEL_CATALOG
-from tests.model_registry.model_catalog.constants import SAMPLE_MODEL_NAME3, CUSTOM_CATALOG_ID1, DEFAULT_CATALOG_ID
-from tests.model_registry.model_catalog.utils import (
+from tests.model_registry.constants import CUSTOM_CATALOG_ID1
+from tests.model_registry.model_catalog.constants import SAMPLE_MODEL_NAME3, DEFAULT_CATALOG_ID
+from tests.model_registry.utils import (
+    get_rest_headers,
     is_model_catalog_ready,
     wait_for_model_catalog_api,
     get_model_str,
     execute_get_command,
+    get_custom_model_catalog_cm_data,
 )
-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_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
-
-
 @pytest.fixture(scope="class")
 def updated_catalog_config_map(
+    pytestconfig: pytest.Config,
     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 = {"data": {"sources.yaml": request.param["sources_yaml"]}}
-    if "sample_yaml" in request.param:
-        for key in request.param["sample_yaml"]:
-            patches["data"][key] = request.param["sample_yaml"][key]
-
-    with ResourceEditor(patches={catalog_config_map: patches}):
+    if pytestconfig.option.post_upgrade or pytestconfig.option.pre_upgrade:
+        yield
+    else:
+        with ResourceEditor(
+            patches={
+                catalog_config_map: get_custom_model_catalog_cm_data(
+                    catalog_config_map=catalog_config_map, param=request.param
+                )
+            }
+        ):
+            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)
-        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)
 
 
 @pytest.fixture(scope="class")

tests/model_registry/model_catalog/constants.py

@@ -1,6 +1,6 @@
-CUSTOM_CATALOG_ID1: str = "sample_custom_catalog1"
+from tests.model_registry.constants import CUSTOM_CATALOG_ID1, SAMPLE_MODEL_NAME1
+
 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 CUSTOM_CATALOG_ID1, SAMPLE_MODEL_NAME1
 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 kubernetes.dynamic.exceptions import ResourceNotFoundError
+from tests.model_registry.utils import (
     get_catalog_str,
     get_sample_yaml_str,
-    ResourceNotFoundError,
+    execute_get_command,
+    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.param(
             {
@@ -62,16 +63,11 @@ 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"]
-
-        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}"
+        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(
         self: Self,

tests/model_registry/model_catalog/test_default_model_catalog.py

@@ -13,12 +13,11 @@
 from tests.model_registry.model_catalog.constants import DEFAULT_CATALOG_ID
 from tests.model_registry.model_catalog.utils import (
     validate_model_catalog_enabled,
-    execute_get_command,
     validate_model_catalog_resource,
     validate_default_catalog,
     get_validate_default_model_catalog_source,
 )
-from tests.model_registry.utils import get_rest_headers
+from tests.model_registry.utils import get_rest_headers, execute_get_command
 from utilities.user_utils import UserTestSession
 
 LOGGER = get_logger(name=__name__)
@@ -32,15 +31,19 @@
 
 @pytest.mark.skip_must_gather
 class TestModelCatalogGeneral:
+    @pytest.mark.pre_upgrade
     @pytest.mark.post_upgrade
     @pytest.mark.install
-    def test_config_map_exists(self: Self, catalog_config_map: ConfigMap):
+    def test_config_map_exists(self: Self, catalog_config_map: ConfigMap, pytestconfig: pytest.Config) -> None:
         # Check that the default configmaps is created when model registry is
         # enabled on data science cluster.
+        expected_num_catalog = 1
+        if pytestconfig.option.pre_upgrade or pytestconfig.option.post_upgrade:
+            expected_num_catalog = 2
         assert catalog_config_map.exists, f"{catalog_config_map.name} does not exist"
         catalogs = yaml.safe_load(catalog_config_map.instance.data["sources.yaml"])["catalogs"]
         assert catalogs
-        assert len(catalogs) == 1, f"{catalog_config_map.name} should have 1 catalog"
+        assert len(catalogs) == expected_num_catalog, f"{catalog_config_map.name} should have 1 catalog"
         validate_default_catalog(default_catalog=catalogs[0])
 
     @pytest.mark.parametrize(
@@ -65,6 +68,7 @@ def test_config_map_exists(self: Self, catalog_config_map: ConfigMap):
         ],
     )
     @pytest.mark.post_upgrade
+    @pytest.mark.pre_upgrade
     @pytest.mark.install
     def test_model_catalog_resources_exists(
         self: Self, admin_client: DynamicClient, model_registry_namespace: str, resource_name: Any
@@ -77,12 +81,14 @@ def test_operator_pod_enabled_model_catalog(self: Self, model_registry_operator_
         assert validate_model_catalog_enabled(pod=model_registry_operator_pod)
 
 
+@pytest.mark.install
 @pytest.mark.parametrize(
     "user_token_for_api_calls,",
     [
         pytest.param(
             {},
             id="test_model_catalog_source_admin_user",
+            marks=(pytest.mark.pre_upgrade, pytest.mark.post_upgrade),
         ),
         pytest.param(
             {"user_type": "test"},
@@ -98,16 +104,28 @@ def test_operator_pod_enabled_model_catalog(self: Self, model_registry_operator_
 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"] == DEFAULT_CATALOG_ID:
+                    items_to_validate.append(catalog)
+        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,