feat: wait for default source state after each test executes (opendatahub-io#1057)

Signed-off-by: lugi0 <lgiorgi@redhat.com>

Author: lugi0

tests/model_registry/model_catalog/catalog_config/test_default_source_inclusion_exclusion_cleanup.py

@@ -16,15 +16,14 @@
     validate_cleanup_logging,
     filter_models_by_pattern,
     execute_inclusion_exclusion_filter_test,
+    ensure_baseline_model_state,
 )
 from tests.model_registry.utils import wait_for_model_catalog_api
 
 LOGGER = get_logger(name=__name__)
 
 pytestmark = [
-    pytest.mark.usefixtures(
-        "updated_dsc_component_state_scope_session", "model_registry_namespace", "baseline_model_state"
-    ),
+    pytest.mark.usefixtures("updated_dsc_component_state_scope_session", "model_registry_namespace"),
 ]
 
 
@@ -65,6 +64,13 @@ def test_include_models_by_pattern(
             model_registry_rest_headers=model_registry_rest_headers,
         )
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
 
 class TestModelExclusionFiltering:
     """Test exclusion filtering functionality (RHOAIENG-41841 part 2)"""
@@ -102,6 +108,13 @@ def test_exclude_models_by_pattern(
             model_registry_rest_headers=model_registry_rest_headers,
         )
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
 
 class TestCombinedIncludeExcludeFiltering:
     """Test combined include+exclude filtering (RHOAIENG-41841 part 3)"""
@@ -191,6 +204,13 @@ def test_combined_include_exclude_filtering(
                 f"SUCCESS: {len(api_models)} {include_pattern} models after excluding {exclude_pattern} variants"
             )
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
 
 class TestModelCleanupLifecycle:
     """Test automatic model cleanup during lifecycle changes (RHOAIENG-41846)"""
@@ -285,6 +305,13 @@ def test_model_cleanup_on_exclusion_change(
                 f"Phase 2 SUCCESS: Granite models cleaned up, {len(phase2_api_models)} prometheus models remain"
             )
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
 
 class TestSourceLifecycleCleanup:
     """Test source disabling cleanup scenarios (RHOAIENG-41846)"""
@@ -330,6 +357,13 @@ def test_source_disabling_removes_models(
 
             LOGGER.info("SUCCESS: Source disabling removed all models")
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
 
 class TestLoggingValidation:
     """Test cleanup operation logging (RHOAIENG-41846)"""
@@ -380,6 +414,13 @@ def test_model_removal_logging(
             except TimeoutExpiredError as e:
                 pytest.fail(f"Expected log patterns not found: {e}")
 
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )
+
     @pytest.mark.sanity
     def test_source_disabling_logging(
         self,
@@ -423,3 +464,10 @@ def test_source_disabling_logging(
                 LOGGER.info(f"SUCCESS: Found expected source disabling log patterns: {found_patterns}")
             except TimeoutExpiredError as e:
                 pytest.fail(f"Expected source disabling log patterns not found: {e}")
+
+        # Ensure baseline model state is restored for subsequent tests
+        ensure_baseline_model_state(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+        )

tests/model_registry/model_catalog/catalog_config/utils.py

@@ -489,3 +489,96 @@ def execute_inclusion_exclusion_filter_test(
         assert api_models == expected_models, f"Expected {test_description}: {expected_models}, got {api_models}"
 
         LOGGER.info(f"SUCCESS: {len(api_models)} {pattern} models {filter_type}")
+
+
+@retry(wait_timeout=300, sleep=10, exceptions_dict={Exception: []}, print_log=False)
+def _validate_baseline_models(
+    model_catalog_rest_url: list[str],
+    model_registry_rest_headers: dict[str, str],
+    model_registry_namespace: str,
+    expected_models: set[str],
+    expected_count: int,
+) -> None:
+    """
+    Validate that baseline model expectations are met.
+    Raises exception if validation fails (triggers retry).
+    Returns None if successful (stops retry).
+    """
+    # Fetch current models from API
+    api_response = get_models_from_catalog_api(
+        model_catalog_rest_url=model_catalog_rest_url,
+        model_registry_rest_headers=model_registry_rest_headers,
+        source_label="Red Hat AI",
+    )
+    api_models = {model["name"] for model in api_response.get("items", [])}
+
+    # Fetch current models from database
+    db_models = get_models_from_database_by_source(source_id=REDHAT_AI_CATALOG_ID, namespace=model_registry_namespace)
+
+    count = len(api_models)
+
+    # Validate all expectations - raise on any failure
+    if count != expected_count:
+        raise AssertionError(f"Expected {expected_count} models, got {count}")
+
+    if api_models != db_models:
+        raise AssertionError(f"API models {api_models} don't match database models {db_models}")
+
+    if api_models != expected_models:
+        raise AssertionError(f"Models {api_models} don't match expected set {expected_models}")
+
+    # Additional category validation
+    granite_models = {model for model in api_models if "granite" in model}
+    prometheus_models = {model for model in api_models if "prometheus" in model}
+
+    if len(granite_models) != 6 or len(prometheus_models) != 1:
+        raise AssertionError(
+            f"""Expected 6 granite + 1 prometheus models, \
+            got {len(granite_models)} granite + {len(prometheus_models)} prometheus"""
+        )
+
+    LOGGER.info("Baseline model validation successful: 7 models (6 granite, 1 prometheus)")
+    return True
+
+
+def ensure_baseline_model_state(
+    model_catalog_rest_url: list[str],
+    model_registry_rest_headers: dict[str, str],
+    model_registry_namespace: str,
+) -> None:
+    """
+    Utility function to ensure that our baseline assumptions about the model data are correct.
+    This should be called at the end of tests to ensure state consistency for subsequent tests.
+    Uses @retry decorator for automatic polling (300s timeout, 10s interval) and eventual reconciliation.
+
+    Args:
+        model_catalog_rest_url: URL for model catalog API
+        model_registry_rest_headers: Headers for API requests
+        model_registry_namespace: Namespace for model registry
+
+    Raises:
+        pytest.FailError: If baseline state cannot be achieved after timeout
+    """
+    # Expected baseline data
+    expected_models = {
+        "granite-3.1-8b-lab-v1",
+        "granite-7b-redhat-lab",
+        "granite-8b-code-base",
+        "granite-8b-code-instruct",
+        "granite-8b-lab-v1",
+        "granite-8b-starter-v1",
+        "prometheus-8x7b-v2-0",
+    }
+    expected_count = 7
+
+    # Use retry decorator for automatic polling and eventual reconciliation
+    try:
+        _validate_baseline_models(
+            model_catalog_rest_url=model_catalog_rest_url,
+            model_registry_rest_headers=model_registry_rest_headers,
+            model_registry_namespace=model_registry_namespace,
+            expected_models=expected_models,
+            expected_count=expected_count,
+        )
+    except TimeoutExpiredError:
+        pytest.fail("Failed to restore baseline model state after 300s timeout")

tests/model_registry/model_catalog/conftest.py

@@ -3,7 +3,6 @@
 import requests
 
 from simple_logger.logger import get_logger
-from timeout_sampler import retry, TimeoutExpiredError
 import yaml
 import pytest
 from kubernetes.dynamic import DynamicClient
@@ -467,95 +466,3 @@ def baseline_redhat_ai_models(
     db_models = get_models_from_database_by_source(source_id=REDHAT_AI_CATALOG_ID, namespace=model_registry_namespace)
 
     return {"api_models": api_models, "db_models": db_models, "count": len(api_models)}
-
-
-@retry(wait_timeout=300, sleep=10, exceptions_dict={Exception: []}, print_log=False)
-def _validate_baseline_models(
-    model_catalog_rest_url: list[str],
-    model_registry_rest_headers: dict[str, str],
-    model_registry_namespace: str,
-    expected_models: set[str],
-    expected_count: int,
-) -> None:
-    """
-    Validate that baseline model expectations are met.
-    Raises exception if validation fails (triggers retry).
-    Returns None if successful (stops retry).
-    """
-    # Fetch current models from API
-    api_response = get_models_from_catalog_api(
-        model_catalog_rest_url=model_catalog_rest_url,
-        model_registry_rest_headers=model_registry_rest_headers,
-        source_label="Red Hat AI",
-    )
-    api_models = {model["name"] for model in api_response.get("items", [])}
-
-    # Fetch current models from database
-    db_models = get_models_from_database_by_source(source_id=REDHAT_AI_CATALOG_ID, namespace=model_registry_namespace)
-
-    count = len(api_models)
-
-    # Validate all expectations - raise on any failure
-    if count != expected_count:
-        raise AssertionError(f"Expected {expected_count} models, got {count}")
-
-    if api_models != db_models:
-        raise AssertionError(f"API models {api_models} don't match database models {db_models}")
-
-    if api_models != expected_models:
-        raise AssertionError(f"Models {api_models} don't match expected set {expected_models}")
-
-    # Additional category validation
-    granite_models = {model for model in api_models if "granite" in model}
-    prometheus_models = {model for model in api_models if "prometheus" in model}
-
-    if len(granite_models) != 6 or len(prometheus_models) != 1:
-        raise AssertionError(
-            f"""Expected 6 granite + 1 prometheus models, \
-            got {len(granite_models)} granite + {len(prometheus_models)} prometheus"""
-        )
-
-    LOGGER.info("Baseline model validation successful: 7 models (6 granite, 1 prometheus)")
-    return True
-
-
-@pytest.fixture(scope="function")
-def baseline_model_state(
-    model_catalog_rest_url: list[str],
-    model_registry_rest_headers: dict[str, str],
-    model_registry_namespace: str,
-) -> None:
-    """
-    Validate that our baseline assumptions about the model data are correct.
-    This fixture should be used by all test classes to ensure data consistency.
-    Uses @retry decorator for automatic polling (300s timeout, 10s interval) and eventual reconciliation.
-
-    Args:
-        model_catalog_rest_url: URL for model catalog API
-        model_registry_rest_headers: Headers for API requests
-        model_registry_namespace: Namespace for model registry
-    """
-
-    # Expected baseline data
-    expected_models = {
-        "granite-3.1-8b-lab-v1",
-        "granite-7b-redhat-lab",
-        "granite-8b-code-base",
-        "granite-8b-code-instruct",
-        "granite-8b-lab-v1",
-        "granite-8b-starter-v1",
-        "prometheus-8x7b-v2-0",
-    }
-    expected_count = 7
-
-    # Use retry decorator for automatic polling and eventual reconciliation
-    try:
-        _validate_baseline_models(
-            model_catalog_rest_url=model_catalog_rest_url,
-            model_registry_rest_headers=model_registry_rest_headers,
-            model_registry_namespace=model_registry_namespace,
-            expected_models=expected_models,
-            expected_count=expected_count,
-        )
-    except TimeoutExpiredError:
-        pytest.fail("Failed to fetch model data after 300s timeout")