New unit tests for jobs with managed clusters

Signed-off-by: Pat O'Connor <[email protected]>

Author: chipspeak

src/codeflare_sdk/ray/cluster/test_cluster.py

@@ -756,6 +756,225 @@ def custom_side_effect(group, version, namespace, plural, **kwargs):
     assert result.dashboard == rc_dashboard
 
 
+def test_run_job_with_managed_cluster_success(mocker):
+    """Test successful RayJob execution with managed cluster."""
+    from codeflare_sdk.ray.job.job import RayJobSpec
+    
+    # Mock Kubernetes API and config
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check")
+    
+    # Mock get_api_client
+    mock_api_client = mocker.Mock()
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client", return_value=mock_api_client)
+    
+    # Mock CustomObjectsApi
+    mock_co_api = mocker.Mock()
+    mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_co_api)
+    
+    # Mock Cluster creation for generating RayCluster spec
+    mock_cluster_resource = {
+        "apiVersion": "ray.io/v1",
+        "kind": "RayCluster",
+        "spec": {
+            "rayVersion": "2.47.1",
+            "headGroupSpec": {"template": {"spec": {}}},
+            "workerGroupSpecs": [{"template": {"spec": {}}}]
+        }
+    }
+    
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster.__init__", return_value=None)
+    mock_cluster_instance = mocker.Mock()
+    mock_cluster_instance.resource_yaml = mock_cluster_resource
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster", return_value=mock_cluster_instance)
+    
+    # Mock successful RayJob creation
+    mock_co_api.create_namespaced_custom_object.return_value = {"metadata": {"name": "test-rayjob"}}
+    
+    # Mock RayJob status for completion
+    mock_status_response = {
+        "status": {
+            "jobDeploymentStatus": "Running",
+            "jobStatus": "SUCCEEDED",
+            "dashboardURL": "http://test-dashboard.com",
+            "rayClusterName": "test-rayjob-cluster",
+            "jobId": "test-job-123"
+        }
+    }
+    mock_co_api.get_namespaced_custom_object_status.return_value = mock_status_response
+    
+    # Mock time.sleep to speed up test
+    mocker.patch("time.sleep")
+    
+    # Create test configuration
+    cluster_config = ClusterConfiguration(
+        name="test-cluster",
+        namespace="test-namespace",
+        num_workers=1,
+        head_cpu_requests=1,
+        head_memory_requests="2G",
+        worker_cpu_requests=1,
+        worker_memory_requests="1G"
+    )
+    
+    job_config = RayJobSpec(
+        entrypoint="python -c 'print(\"Hello World\")'",
+        submission_id="test-submission-123"
+    )
+    
+    # Execute the method
+    result = Cluster.run_job_with_managed_cluster(
+        cluster_config=cluster_config,
+        job_config=job_config,
+        job_cr_name="test-rayjob",
+        wait_for_completion=True,
+        job_timeout_seconds=60
+    )
+    
+    # Verify the result
+    assert result["job_cr_name"] == "test-rayjob"
+    assert result["job_submission_id"] == "test-job-123"
+    assert result["job_status"] == "SUCCEEDED"
+    assert result["dashboard_url"] == "http://test-dashboard.com"
+    assert result["ray_cluster_name"] == "test-rayjob-cluster"
+    
+    # Verify API calls were made
+    mock_co_api.create_namespaced_custom_object.assert_called_once()
+    mock_co_api.get_namespaced_custom_object_status.assert_called()
+
+
+def test_run_job_with_managed_cluster_no_wait(mocker):
+    """Test RayJob execution without waiting for completion."""
+    from codeflare_sdk.ray.job.job import RayJobSpec
+    
+    # Mock dependencies (similar to above but condensed)
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check")
+    
+    mock_api_client = mocker.Mock()
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client", return_value=mock_api_client)
+    
+    mock_co_api = mocker.Mock()
+    mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_co_api)
+    
+    # Mock Cluster creation
+    mock_cluster_resource = {
+        "apiVersion": "ray.io/v1",
+        "kind": "RayCluster", 
+        "spec": {"rayVersion": "2.47.1", "headGroupSpec": {"template": {"spec": {}}}}
+    }
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster.__init__", return_value=None)
+    mock_cluster_instance = mocker.Mock()
+    mock_cluster_instance.resource_yaml = mock_cluster_resource
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster", return_value=mock_cluster_instance)
+    
+    # Mock RayJob creation and status (not found initially - job just submitted)
+    mock_co_api.create_namespaced_custom_object.return_value = {"metadata": {"name": "test-rayjob-nowait"}}
+    
+    from kubernetes.client.rest import ApiException
+    mock_co_api.get_namespaced_custom_object_status.side_effect = ApiException(status=404, reason="Not Found")
+    
+    cluster_config = ClusterConfiguration(
+        name="test-cluster",
+        namespace="test-namespace"
+    )
+    
+    job_config = RayJobSpec(entrypoint="python script.py")
+    
+    # Execute without waiting
+    result = Cluster.run_job_with_managed_cluster(
+        cluster_config=cluster_config,
+        job_config=job_config,
+        wait_for_completion=False
+    )
+    
+    # Verify result for no-wait case
+    assert "job_cr_name" in result
+    assert result["job_status"] == "SUBMITTED_NOT_FOUND"
+    
+    # Verify no polling happened
+    mock_co_api.create_namespaced_custom_object.assert_called_once()
+
+
+def test_run_job_with_managed_cluster_timeout(mocker):
+    """Test RayJob execution with timeout."""
+    from codeflare_sdk.ray.job.job import RayJobSpec
+    
+    # Mock dependencies
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.config_check")
+    
+    mock_api_client = mocker.Mock()
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client", return_value=mock_api_client)
+    
+    mock_co_api = mocker.Mock()
+    mocker.patch("kubernetes.client.CustomObjectsApi", return_value=mock_co_api)
+    
+    # Mock Cluster creation
+    mock_cluster_resource = {"apiVersion": "ray.io/v1", "kind": "RayCluster", "spec": {}}
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster.__init__", return_value=None)
+    mock_cluster_instance = mocker.Mock()
+    mock_cluster_instance.resource_yaml = mock_cluster_resource
+    mocker.patch("codeflare_sdk.ray.cluster.cluster.Cluster", return_value=mock_cluster_instance)
+    
+    # Mock RayJob creation
+    mock_co_api.create_namespaced_custom_object.return_value = {"metadata": {"name": "test-rayjob-timeout"}}
+    
+    # Mock job status as always RUNNING (never completes)
+    mock_status_response = {
+        "status": {
+            "jobDeploymentStatus": "Running", 
+            "jobStatus": "RUNNING",
+            "jobId": "timeout-job-123"
+        }
+    }
+    mock_co_api.get_namespaced_custom_object_status.return_value = mock_status_response
+    
+    # Mock time to simulate timeout
+    start_time = 1000
+    mocker.patch("time.time", side_effect=[start_time, start_time + 70])  # Exceed 60s timeout
+    mocker.patch("time.sleep")
+    
+    cluster_config = ClusterConfiguration(name="timeout-cluster", namespace="test-ns")
+    job_config = RayJobSpec(entrypoint="python long_running_script.py")
+    
+    # Execute and expect timeout
+    try:
+        result = Cluster.run_job_with_managed_cluster(
+            cluster_config=cluster_config,
+            job_config=job_config,
+            wait_for_completion=True,
+            job_timeout_seconds=60,
+            job_polling_interval_seconds=1
+        )
+        assert False, "Expected TimeoutError"
+    except TimeoutError as e:
+        assert "timed out after 60 seconds" in str(e)
+        assert "RUNNING" in str(e)
+
+
+def test_run_job_with_managed_cluster_validation_error(mocker):
+    """Test RayJob execution with validation errors."""
+    from codeflare_sdk.ray.job.job import RayJobSpec
+    
+    cluster_config = ClusterConfiguration(name="test-cluster", namespace="test-ns")
+    
+    # Test missing entrypoint
+    job_config_no_entrypoint = RayJobSpec(entrypoint="")
+    
+    try:
+        Cluster.run_job_with_managed_cluster(
+            cluster_config=cluster_config,
+            job_config=job_config_no_entrypoint
+        )
+        assert False, "Expected ValueError for missing entrypoint"
+    except ValueError as e:
+        assert "entrypoint must be specified" in str(e)
+
+
 # Make sure to always keep this function last
 def test_cleanup():
     os.remove(f"{aw_dir}test-all-params.yaml")