RHOAIENG-32532: Add additional test coverage

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -2004,3 +2004,193 @@ def test_update_existing_cluster_for_scripts_api_errors(mocker, auto_mock_setup)
         rayjob._update_existing_cluster_for_scripts(
             "test-configmap", mock_config_builder
         )
+
+
+def test_rayjob_kueue_label_no_default_queue(auto_mock_setup, mocker, caplog):
+    """Test RayJob falls back to 'default' queue when no default queue exists."""
+    mocker.patch(
+        "codeflare_sdk.ray.rayjobs.rayjob.get_default_kueue_name",
+        return_value=None,
+    )
+
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+    mock_api_instance.submit_job.return_value = {"metadata": {"name": "test-job"}}
+
+    cluster_config = ManagedClusterConfig()
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_config=cluster_config,
+        entrypoint="python script.py",
+    )
+
+    with caplog.at_level("WARNING"):
+        rayjob.submit()
+
+    # Verify the submitted job has the fallback label
+    call_args = mock_api_instance.submit_job.call_args
+    submitted_job = call_args.kwargs["job"]
+    assert submitted_job["metadata"]["labels"]["kueue.x-k8s.io/queue-name"] == "default"
+
+    # Verify warning was logged
+    assert "No default Kueue LocalQueue found" in caplog.text
+
+
+def test_rayjob_kueue_explicit_local_queue(auto_mock_setup):
+    """Test RayJob uses explicitly specified local queue."""
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+    mock_api_instance.submit_job.return_value = {"metadata": {"name": "test-job"}}
+
+    cluster_config = ManagedClusterConfig()
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_config=cluster_config,
+        entrypoint="python script.py",
+        local_queue="custom-queue",
+    )
+
+    rayjob.submit()
+
+    # Verify the submitted job has the explicit queue label
+    call_args = mock_api_instance.submit_job.call_args
+    submitted_job = call_args.kwargs["job"]
+    assert (
+        submitted_job["metadata"]["labels"]["kueue.x-k8s.io/queue-name"]
+        == "custom-queue"
+    )
+
+
+def test_rayjob_no_kueue_label_for_existing_cluster(auto_mock_setup):
+    """Test RayJob doesn't add Kueue label for existing clusters."""
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+    mock_api_instance.submit_job.return_value = {"metadata": {"name": "test-job"}}
+
+    # Using existing cluster (no cluster_config)
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_name="existing-cluster",
+        entrypoint="python script.py",
+    )
+
+    rayjob.submit()
+
+    # Verify no Kueue label was added
+    call_args = mock_api_instance.submit_job.call_args
+    submitted_job = call_args.kwargs["job"]
+    assert "kueue.x-k8s.io/queue-name" not in submitted_job["metadata"]["labels"]
+
+
+def test_rayjob_with_ttl_and_deadline(auto_mock_setup):
+    """Test RayJob with TTL and active deadline seconds."""
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+    mock_api_instance.submit_job.return_value = {"metadata": {"name": "test-job"}}
+
+    cluster_config = ManagedClusterConfig()
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_config=cluster_config,
+        entrypoint="python script.py",
+        ttl_seconds_after_finished=300,
+        active_deadline_seconds=600,
+    )
+
+    rayjob.submit()
+
+    # Verify TTL and deadline were set
+    call_args = mock_api_instance.submit_job.call_args
+    submitted_job = call_args.kwargs["job"]
+    assert submitted_job["spec"]["ttlSecondsAfterFinished"] == 300
+    assert submitted_job["spec"]["activeDeadlineSeconds"] == 600
+
+
+def test_rayjob_shutdown_after_job_finishes(auto_mock_setup):
+    """Test RayJob sets shutdownAfterJobFinishes correctly."""
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+    mock_api_instance.submit_job.return_value = {"metadata": {"name": "test-job"}}
+
+    # Test with managed cluster (should shutdown)
+    cluster_config = ManagedClusterConfig()
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_config=cluster_config,
+        entrypoint="python script.py",
+    )
+
+    rayjob.submit()
+
+    call_args = mock_api_instance.submit_job.call_args
+    submitted_job = call_args.kwargs["job"]
+    assert submitted_job["spec"]["shutdownAfterJobFinishes"] is True
+
+    # Test with existing cluster (should not shutdown)
+    rayjob2 = RayJob(
+        job_name="test-job2",
+        cluster_name="existing-cluster",
+        entrypoint="python script.py",
+    )
+
+    rayjob2.submit()
+
+    call_args2 = mock_api_instance.submit_job.call_args
+    submitted_job2 = call_args2.kwargs["job"]
+    assert submitted_job2["spec"]["shutdownAfterJobFinishes"] is False
+
+
+def test_rayjob_stop_delete_resubmit_logging(auto_mock_setup, caplog):
+    """Test logging for stop, delete, and resubmit operations."""
+    mock_api_instance = auto_mock_setup["rayjob_api"]
+
+    # Test stop with logging
+    mock_api_instance.suspend_job.return_value = {
+        "metadata": {"name": "test-rayjob"},
+        "spec": {"suspend": True},
+    }
+
+    rayjob = RayJob(
+        job_name="test-rayjob",
+        cluster_name="test-cluster",
+        namespace="test-namespace",
+        entrypoint="python script.py",
+    )
+
+    with caplog.at_level("INFO"):
+        result = rayjob.stop()
+
+    assert result is True
+    assert "Successfully stopped the RayJob test-rayjob" in caplog.text
+
+    # Test delete with logging
+    caplog.clear()
+    mock_api_instance.delete_job.return_value = True
+
+    with caplog.at_level("INFO"):
+        result = rayjob.delete()
+
+    assert result is True
+    assert "Successfully deleted the RayJob test-rayjob" in caplog.text
+
+    # Test resubmit with logging
+    caplog.clear()
+    mock_api_instance.resubmit_job.return_value = {
+        "metadata": {"name": "test-rayjob"},
+        "spec": {"suspend": False},
+    }
+
+    with caplog.at_level("INFO"):
+        result = rayjob.resubmit()
+
+    assert result is True
+    assert "Successfully resubmitted the RayJob test-rayjob" in caplog.text
+
+
+def test_rayjob_initialization_logging(auto_mock_setup, caplog):
+    """Test RayJob initialization logging."""
+    with caplog.at_level("INFO"):
+        cluster_config = ManagedClusterConfig()
+        rayjob = RayJob(
+            job_name="test-job",
+            cluster_config=cluster_config,
+            entrypoint="python script.py",
+        )
+
+    assert "Creating new cluster: test-job-cluster" in caplog.text
+    assert "Initialized RayJob: test-job in namespace: test-namespace" in caplog.text

src/codeflare_sdk/ray/rayjobs/test_status.py

@@ -302,3 +302,72 @@ def test_rayjob_status_print_job_found(mocker):
 
     assert status == CodeflareRayJobStatus.RUNNING
     assert ready == False
+
+
+def test_rayjob_status_all_deployment_states(mocker):
+    """Test RayJob status method with all deployment states."""
+    mocker.patch("kubernetes.config.load_kube_config")
+    mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+    mock_api_instance = mock_api_class.return_value
+
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_name="test-cluster",
+        namespace="test-ns",
+        entrypoint="python test.py",
+    )
+
+    # Test mapping of deployment statuses to CodeflareRayJobStatus
+    test_cases = [
+        # (deployment_status_str, expected CodeflareRayJobStatus, expected ready)
+        ("Complete", CodeflareRayJobStatus.COMPLETE, True),
+        ("Running", CodeflareRayJobStatus.RUNNING, False),
+        ("Failed", CodeflareRayJobStatus.FAILED, False),
+        ("Suspended", CodeflareRayJobStatus.SUSPENDED, False),
+    ]
+
+    for deployment_status_str, expected_status, expected_ready in test_cases:
+        mock_api_instance.get_job_status.return_value = {
+            "jobId": "test-job-abc123",
+            "jobDeploymentStatus": deployment_status_str,
+            "startTime": "2025-07-28T11:37:07Z",
+            "failed": 0,
+            "succeeded": 0,
+            "rayClusterName": "test-cluster",
+        }
+        status, ready = rayjob.status(print_to_console=False)
+        assert status == expected_status, f"Failed for {deployment_status_str}"
+        assert (
+            ready == expected_ready
+        ), f"Failed ready check for {deployment_status_str}"
+
+
+def test_rayjob_status_with_end_time(mocker):
+    """Test RayJob status with end time field."""
+    mocker.patch("kubernetes.config.load_kube_config")
+    mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
+    mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayClusterApi")
+    mock_api_instance = mock_api_class.return_value
+
+    rayjob = RayJob(
+        job_name="test-job",
+        cluster_name="test-cluster",
+        namespace="test-ns",
+        entrypoint="python test.py",
+    )
+
+    # Test with end time field
+    mock_api_instance.get_job_status.return_value = {
+        "jobId": "test-job-abc123",
+        "jobDeploymentStatus": "Complete",
+        "startTime": "2025-07-28T11:37:07Z",
+        "endTime": "2025-07-28T11:47:07Z",
+        "failed": 0,
+        "succeeded": 1,
+        "rayClusterName": "test-cluster",
+    }
+
+    status, ready = rayjob.status(print_to_console=False)
+    assert status == CodeflareRayJobStatus.COMPLETE
+    assert ready == True

tests/e2e/rayjob/rayjob_lifecycled_cluster_test.py

@@ -120,9 +120,28 @@ def test_lifecycled_kueue_resource_queueing(self):
             )
             assert job2.submit() == "waiter"
 
-            sleep(3)
+            # Wait for Kueue to process the job
+            sleep(5)
             job2_cr = self.job_api.get_job(name=job2.name, k8s_namespace=job2.namespace)
-            assert job2_cr.get("spec", {}).get("suspend", False)
+
+            # For RayJobs with managed clusters, check if Kueue is holding resources
+            job2_status = job2_cr.get("status", {})
+            ray_cluster_name = job2_status.get("rayClusterName", "")
+
+            # If RayCluster is not created yet, it means Kueue is holding the job
+            if not ray_cluster_name:
+                # This is the expected behavior
+                job_is_queued = True
+            else:
+                # Check RayCluster resources - if all are 0, it's queued
+                ray_cluster_status = job2_status.get("rayClusterStatus", {})
+                desired_cpu = ray_cluster_status.get("desiredCPU", "0")
+                desired_memory = ray_cluster_status.get("desiredMemory", "0")
+
+                # Kueue creates the RayCluster but with 0 resources when queued
+                job_is_queued = desired_cpu == "0" and desired_memory == "0"
+
+            assert job_is_queued, "Job2 should be queued by Kueue while Job1 is running"
 
             assert self.job_api.wait_until_job_finished(
                 name=job1.name, k8s_namespace=job1.namespace, timeout=60

tests/e2e/support.py

@@ -559,7 +559,7 @@ def create_limited_kueue_resources(self):
     if is_openshift():
         # MODH images need more memory, so higher quota but still limited to allow only 1 job
         cpu_quota = 3
-        memory_quota = "15Gi"
+        memory_quota = "8Gi"  # Reduced from 15Gi to ensure second job is suspended (1 job needs 7Gi)
     else:
         # Standard Ray images need less memory
         cpu_quota = 3