refactor the collect_pod_information function to be re-used in different component tests

Signed-off-by: Kamesh Akella <kakella@redhat.com>

Author: kami619

tests/model_registry/async_job/utils.py

@@ -10,6 +10,7 @@
 from timeout_sampler import TimeoutExpiredError
 
 from utilities.must_gather_collector import get_base_dir, get_must_gather_collector_dir
+from utilities.general import collect_pod_information
 
 LOGGER = get_logger(name=__name__)
 
@@ -136,22 +137,3 @@ def upload_test_model_to_minio_from_image(
         LOGGER.info(
             f"Test model file uploaded successfully to s3://{MinIo.Buckets.MODELMESH_EXAMPLE_MODELS}/{object_key}"
         )
-
-
-def collect_pod_information(pod: Pod) -> None:
-    try:
-        base_dir_name = get_must_gather_collector_dir() or get_base_dir()
-        LOGGER.info(f"Collecting pod information for {pod.name}: {base_dir_name}")
-        os.makedirs(base_dir_name, exist_ok=True)
-        yaml_file_path = os.path.join(base_dir_name, f"{pod.name}.yaml")
-        with open(yaml_file_path, "w") as fd:
-            fd.write(pod.instance.to_str())
-        # get all the containers of the pod:
-
-        containers = [container["name"] for container in pod.instance.status.containerStatuses]
-        for container in containers:
-            file_path = os.path.join(base_dir_name, f"{pod.name}_{container}.log")
-            with open(file_path, "w") as fd:
-                fd.write(pod.log(**{"container": container}))
-    except Exception:
-        LOGGER.warning(f"For pod: {pod.name} information gathering failed.")

tests/workbenches/conftest.py

@@ -17,6 +17,7 @@
 from utilities import constants
 from utilities.constants import INTERNAL_IMAGE_REGISTRY_PATH
 from utilities.infra import check_internal_image_registry_available
+from utilities.general import collect_pod_information
 
 LOGGER = get_logger(name=__name__)
 
@@ -67,16 +68,21 @@ def notebook_image(
         if not custom_image:
             raise ValueError("custom_image cannot be empty or whitespace")
 
-        # Validation Logic (Moved from default_notebook)
+        # Validation Logic: Only digest references are accepted
         _ERR_INVALID_CUSTOM_IMAGE = (
-            "custom_image must be a valid OCI image reference with either a tag (:tag) or digest (@sha256:digest), "
-            "e.g., 'quay.io/org/image:tag' or 'quay.io/org/image@sha256:digest', "
+            "custom_image must be a valid OCI image reference with a digest (@sha256:digest), "
+            "e.g., 'quay.io/org/image@sha256:abc123...', "
             "got: '{custom_image}'"
         )
-        has_digest = "@sha256:" in custom_image
-        has_tag = ":" in custom_image and custom_image.rfind(":") > custom_image.rfind("/")
-
-        if not (has_digest or has_tag):
+        # Check for valid digest: @sha256: must be followed by non-empty content
+        digest_marker = "@sha256:"
+        has_valid_digest = False
+        if digest_marker in custom_image:
+            digest_index = custom_image.rfind(digest_marker)
+            digest_end = digest_index + len(digest_marker)
+            has_valid_digest = digest_end < len(custom_image)
+
+        if not has_valid_digest:
             raise ValueError(_ERR_INVALID_CUSTOM_IMAGE.format(custom_image=custom_image))
 
         LOGGER.info(f"Using custom workbench image: {custom_image}")
@@ -261,6 +267,8 @@ def notebook_pod(
         )
     except (TimeoutError, RuntimeError) as e:
         if notebook_pod.exists:
+            # Collect pod information for debugging purposes
+            collect_pod_information(notebook_pod)
             pod_status = notebook_pod.instance.status
             pod_phase = pod_status.phase
             error_details = get_pod_failure_details(notebook_pod)

utilities/general.py

@@ -2,6 +2,7 @@
 import re
 from typing import List, Tuple, Any
 import uuid
+import os
 
 from kubernetes.dynamic import DynamicClient
 from kubernetes.dynamic.exceptions import ResourceNotFoundError, NotFoundError
@@ -13,6 +14,7 @@
 import utilities.infra
 from utilities.constants import Annotations, KServeDeploymentType, MODELMESH_SERVING
 from utilities.exceptions import UnexpectedResourceCountError, ResourceValueMismatch
+from utilities.must_gather_collector import get_base_dir, get_must_gather_collector_dir
 from ocp_resources.resource import Resource
 from timeout_sampler import retry
 from timeout_sampler import TimeoutExpiredError, TimeoutSampler
@@ -467,3 +469,22 @@ def wait_for_pods_running(
             )
             raise
     return None
+
+
+def collect_pod_information(pod: Pod) -> None:
+    try:
+        base_dir_name = get_must_gather_collector_dir() or get_base_dir()
+        LOGGER.info(f"Collecting pod information for {pod.name}: {base_dir_name}")
+        os.makedirs(base_dir_name, exist_ok=True)
+        yaml_file_path = os.path.join(base_dir_name, f"{pod.name}.yaml")
+        with open(yaml_file_path, "w") as fd:
+            fd.write(pod.instance.to_str())
+        # get all the containers of the pod:
+
+        containers = [container["name"] for container in pod.instance.status.containerStatuses]
+        for container in containers:
+            file_path = os.path.join(base_dir_name, f"{pod.name}_{container}.log")
+            with open(file_path, "w") as fd:
+                fd.write(pod.log(**{"container": container}))
+    except Exception:
+        LOGGER.warning(f"For pod: {pod.name} information gathering failed.")