address comments

Author: dbasunag

tests/model_registry/model_catalog/huggingface/conftest.py

@@ -2,7 +2,8 @@
 import time
 from huggingface_hub import HfApi
 from simple_logger.logger import get_logger
-from tests.model_registry.utils import execute_get_command
+
+from tests.model_registry.model_catalog.huggingface.utils import get_huggingface_model_from_api
 
 LOGGER = get_logger(name=__name__)
 
@@ -56,11 +57,11 @@ def initial_last_synced_values(
     """
     Collect initial last_synced values for a given model.
     """
-    model_name = request.param
-    url = f"{model_catalog_rest_url[0]}sources/hf_id/models/{model_name}"
-    result = execute_get_command(
-        url=url,
-        headers=model_registry_rest_headers,
+    result = get_huggingface_model_from_api(
+        model_registry_rest_headers=model_registry_rest_headers,
+        model_catalog_rest_url=model_catalog_rest_url,
+        model_name=request.param,
+        source_id="hf_id",
     )
 
     return result["customProperties"]["last_synced"]["string_value"]

tests/model_registry/model_catalog/huggingface/test_huggingface_model_validation.py

@@ -6,12 +6,12 @@
 from tests.model_registry.model_catalog.utils import (
     get_hf_catalog_str,
 )
-from tests.model_registry.utils import execute_get_command
 from tests.model_registry.model_catalog.huggingface.utils import (
     assert_huggingface_values_matches_model_catalog_api_values,
     wait_for_huggingface_retrival_match,
     wait_for_hugging_face_model_import,
     wait_for_last_sync_update,
+    get_huggingface_model_from_api,
 )
 from huggingface_hub import HfApi
 from kubernetes.dynamic import DynamicClient
@@ -60,6 +60,7 @@ def test_huggingface_last_synced_custom(
             model_registry_rest_headers=model_registry_rest_headers,
             model_catalog_rest_url=model_catalog_rest_url,
             model_name=model_name,
+            source_id="hf_id",
             initial_last_synced_values=float(initial_last_synced_values),
         )
 
@@ -85,7 +86,6 @@ class TestHuggingFaceModelValidation:
     def test_huggingface_model_metadata_last_synced(
         self: Self,
         epoch_time_before_config_map_update: float,
-        updated_catalog_config_map: tuple[ConfigMap, str, str],
         model_catalog_rest_url: list[str],
         model_registry_rest_headers: dict[str, str],
         expected_catalog_values: dict[str, str],
@@ -100,12 +100,12 @@ def test_huggingface_model_metadata_last_synced(
         )
         error = {}
         for model_name in expected_catalog_values:
-            url = f"{model_catalog_rest_url[0]}sources/{HF_SOURCE_ID}/models/{model_name}"
-            result = execute_get_command(
-                url=url,
-                headers=model_registry_rest_headers,
+            result = get_huggingface_model_from_api(
+                model_catalog_rest_url=model_catalog_rest_url,
+                model_registry_rest_headers=model_registry_rest_headers,
+                model_name=model_name,
+                source_id=HF_SOURCE_ID,
             )
-
             error_msg = ""
             if result["name"] != model_name:
                 error_msg += f"Expected model name {model_name}, but got {result['name']}. "
@@ -115,17 +115,13 @@ def test_huggingface_model_metadata_last_synced(
             LOGGER.info(f"Model {model_name} last synced at: {last_synced}")
 
             # Validate that last_synced field exists and is not empty
-            if last_synced is None:
-                error_msg += f"last_synced field is None for model {model_name}. "
-            elif last_synced == "":
-                error_msg += f"last_synced field is empty for model {model_name}. "
-            else:
-                # Compare timestamps: current_epoch_time should be earlier than last_synced
-                if epoch_time_before_config_map_update > float(last_synced):
-                    error_msg += (
-                        f"Model {model_name} last_synced ({last_synced}) should be after "
-                        f"test start time ({epoch_time_before_config_map_update}). "
-                    )
+            if not last_synced or last_synced == "":
+                error_msg += f"last_synced field is not present for model {model_name}. "
+            elif epoch_time_before_config_map_update > float(last_synced):
+                error_msg += (
+                    f"Model {model_name} last_synced ({last_synced}) should be after "
+                    f"test start time ({epoch_time_before_config_map_update}). "
+                )
             if error_msg:
                 error[model_name] = error_msg
         if error:

tests/model_registry/model_catalog/huggingface/utils.py

@@ -1,5 +1,6 @@
 import ast
 from typing import Any
+
 from simple_logger.logger import get_logger
 
 from tests.model_registry.model_catalog.constants import HF_SOURCE_ID
@@ -142,20 +143,35 @@ def wait_for_hugging_face_model_import(
         return False
 
 
+def get_huggingface_model_from_api(
+    model_catalog_rest_url: list[str],
+    model_registry_rest_headers: dict[str, str],
+    model_name: str,
+    source_id: str,
+) -> dict[str, Any]:
+    url = f"{model_catalog_rest_url[0]}sources/{source_id}/models/{model_name}"
+    return execute_get_command(
+        url=url,
+        headers=model_registry_rest_headers,
+    )
+
+
 @retry(wait_timeout=135, sleep=15)
 def wait_for_last_sync_update(
     model_catalog_rest_url: list[str],
     model_registry_rest_headers: dict[str, str],
     model_name: str,
+    source_id: str,
     initial_last_synced_values: float,
 ) -> bool:
     """Wait for the last_synced value to be updated with exact 120-second difference"""
-    url = f"{model_catalog_rest_url[0]}sources/hf_id/models/{model_name}"
-    result = execute_get_command(
-        url=url,
-        headers=model_registry_rest_headers,
-    )
 
+    result = get_huggingface_model_from_api(
+        model_registry_rest_headers=model_registry_rest_headers,
+        model_catalog_rest_url=model_catalog_rest_url,
+        model_name=model_name,
+        source_id=source_id,
+    )
     current_last_synced = float(result["customProperties"]["last_synced"]["string_value"])
     if current_last_synced != initial_last_synced_values:
         # Calculate difference in milliseconds and convert to seconds