opendatahub-io
diff --git a/‎tests/model_registry/async_job/utils.py‎
Lines changed: 1 addition & 22 deletions b/‎tests/model_registry/async_job/utils.py‎
Lines changed: 1 addition & 22 deletions
diff --git a/‎tests/workbenches/conftest.py‎
Lines changed: 130 additions & 14 deletions b/‎tests/workbenches/conftest.py‎
Lines changed: 130 additions & 14 deletions
@@ -1,5 +1,3 @@
-import os
-
 import requests
 from kubernetes.dynamic import DynamicClient
 from ocp_resources.job import Job
@@ -9,7 +7,7 @@
 from simple_logger.logger import get_logger
 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 +134,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.")
@@ -11,11 +11,13 @@
 from ocp_resources.namespace import Namespace
 from ocp_resources.persistent_volume_claim import PersistentVolumeClaim
 from ocp_resources.notebook import Notebook
+from ocp_resources.pod import Pod
 
-from utilities.constants import Labels
+from utilities.constants import Labels, Timeout
 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__)
 
@@ -44,10 +46,64 @@ def minimal_image() -> Generator[str, None, None]:
 
 
 @pytest.fixture(scope="function")
-def default_notebook(
+def notebook_image(
     request: pytest.FixtureRequest,
     admin_client: DynamicClient,
     minimal_image: str,
+) -> str:
+    """
+    Resolves the notebook image path.
+
+    Priority:
+    1. 'custom_image' provided via indirect parametrization
+    2. Default 'minimal_image' (with automatic registry resolution)
+    """
+    # SAFELY get parameters. If test doesn't parameterize this fixture, default to empty dict.
+    params = getattr(request, "param", {})
+    custom_image = params.get("custom_image")
+
+    # Case A: Custom Image (Explicit)
+    if custom_image:
+        custom_image = custom_image.strip()
+        if not custom_image:
+            raise ValueError("custom_image cannot be empty or whitespace")
+
+        # Validation Logic: Only digest references are accepted
+        _ERR_INVALID_CUSTOM_IMAGE = (
+            "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}'"
+        )
+        # 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}")
+        return custom_image
+
+    # Case B: Default Image (Implicit / Good Default)
+    # This runs for all standard tests in test_spawning.py
+    internal_image_registry = check_internal_image_registry_available(admin_client=admin_client)
+
+    return (
+        f"{INTERNAL_IMAGE_REGISTRY_PATH}/{py_config['applications_namespace']}/{minimal_image}"
+        if internal_image_registry
+        else minimal_image
+    )
+
+
+@pytest.fixture(scope="function")
+def default_notebook(
+    request: pytest.FixtureRequest,
+    admin_client: DynamicClient,
+    notebook_image: str,
 ) -> Generator[Notebook, None, None]:
     """Returns a new Notebook CR for a given namespace, name, and image"""
     namespace = request.param["namespace"]
@@ -60,15 +116,8 @@ def default_notebook(
     username = get_username(dyn_client=admin_client)
     assert username, "Failed to determine username from the cluster"
 
-    # Check internal image registry availability
-    internal_image_registry = check_internal_image_registry_available(admin_client=admin_client)
-
-    # Set the image path based on internal image registry status
-    minimal_image_path = (
-        f"{INTERNAL_IMAGE_REGISTRY_PATH}/{py_config['applications_namespace']}/{minimal_image}"
-        if internal_image_registry
-        else ":" + minimal_image.rsplit(":", maxsplit=1)[1]
-    )
+    # Set the image path based on the resolved notebook_image
+    image_path = notebook_image
 
     probe_config = {
         "failureThreshold": 3,
@@ -90,7 +139,7 @@ def default_notebook(
             "annotations": {
                 Labels.Notebook.INJECT_AUTH: "true",
                 "opendatahub.io/accelerator-name": "",
-                "notebooks.opendatahub.io/last-image-selection": minimal_image,
+                "notebooks.opendatahub.io/last-image-selection": image_path,
                 # Add any additional annotations if provided
                 **auth_annotations,
             },
@@ -124,9 +173,9 @@ def default_notebook(
                                     "                  "
                                     "--ServerApp.quit_button=False\n",
                                 },
-                                {"name": "JUPYTER_IMAGE", "value": minimal_image_path},
+                                {"name": "JUPYTER_IMAGE", "value": image_path},
                             ],
-                            "image": minimal_image_path,
+                            "image": image_path,
                             "imagePullPolicy": "Always",
                             "livenessProbe": probe_config,
                             "name": name,
@@ -167,3 +216,70 @@ def default_notebook(
 
     with Notebook(kind_dict=notebook) as nb:
         yield nb
+
+
+@pytest.fixture(scope="function")
+def notebook_pod(
+    unprivileged_client: DynamicClient,
+    default_notebook: Notebook,
+) -> Pod:
+    """
+    Returns a notebook pod in Ready state.
+
+    This fixture:
+    - Creates a Pod object for the notebook
+    - Waits for pod to exist
+    - Waits for pod to reach Ready state (10-minute timeout)
+    - Provides detailed diagnostics on failure
+
+    Args:
+        unprivileged_client: Client for interacting with the cluster
+        default_notebook: The notebook CR to get the pod for
+
+    Returns:
+        Pod object in Ready state
+
+    Raises:
+        AssertionError: If pod fails to reach Ready state or is not created
+    """
+    # Error messages
+    _ERR_POD_NOT_READY = (
+        "Pod '{pod_name}-0' failed to reach Ready state within 10 minutes.\n"
+        "Pod Phase: {pod_phase}\n"
+        "Original Error: {original_error}\n"
+        "Pod information collected to must-gather directory for debugging."
+    )
+    _ERR_POD_NOT_CREATED = "Pod '{pod_name}-0' was not created. Check notebook controller logs."
+
+    # Create pod object
+    notebook_pod = Pod(
+        client=unprivileged_client,
+        namespace=default_notebook.namespace,
+        name=f"{default_notebook.name}-0",
+    )
+
+    try:
+        notebook_pod.wait()
+        notebook_pod.wait_for_condition(
+            condition=Pod.Condition.READY,
+            status=Pod.Condition.Status.TRUE,
+            timeout=Timeout.TIMEOUT_10MIN,
+        )
+    except (TimeoutError, RuntimeError) as e:
+        if notebook_pod.exists:
+            # Collect pod information for debugging purposes (YAML + logs saved to must-gather dir)
+            collect_pod_information(notebook_pod)
+            pod_status = notebook_pod.instance.status
+            pod_phase = pod_status.phase
+            raise AssertionError(
+                _ERR_POD_NOT_READY.format(
+                    pod_name=default_notebook.name,
+                    pod_phase=pod_phase,
+                    original_error=e,
+                )
+            ) from e
+        else:
+            # Pod was never created
+            raise AssertionError(_ERR_POD_NOT_CREATED.format(pod_name=default_notebook.name)) from e
+
+    return notebook_pod