Trying to bump job file coverage

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

Author: chipspeak

src/codeflare_sdk/ray/job/test_job.py

@@ -0,0 +1,225 @@
+import pytest
+from codeflare_sdk.ray.job.job import RayJobSpec, RayJob, RayJobStatus
+
+
+def test_ray_job_spec_creation():
+    """Test RayJobSpec creation with all fields."""
+    spec = RayJobSpec(
+        entrypoint="python script.py",
+        submission_id="test-123",
+        runtime_env={"pip": ["numpy", "pandas"]},
+        metadata={"author": "test"},
+        entrypoint_num_cpus=2.0,
+        entrypoint_num_gpus=1.0,
+        entrypoint_memory=1024,
+        entrypoint_resources={"custom.com/resource": 1.0},
+        cluster_name="test-cluster",
+        cluster_namespace="test-ns"
+    )
+    
+    assert spec.entrypoint == "python script.py"
+    assert spec.submission_id == "test-123"
+    assert spec.runtime_env == {"pip": ["numpy", "pandas"]}
+    assert spec.metadata == {"author": "test"}
+    assert spec.entrypoint_num_cpus == 2.0
+    assert spec.entrypoint_num_gpus == 1.0
+    assert spec.entrypoint_memory == 1024
+    assert spec.entrypoint_resources == {"custom.com/resource": 1.0}
+    assert spec.cluster_name == "test-cluster"
+    assert spec.cluster_namespace == "test-ns"
+    assert spec.status == RayJobStatus.PENDING
+
+
+def test_ray_job_yaml_generation_full():
+    """Test RayJob YAML generation with all fields."""
+    spec = RayJobSpec(
+        entrypoint="python -c 'import ray; print(ray.cluster_resources())'",
+        submission_id="test-submission-456",
+        runtime_env={"pip": ["numpy==1.24.0", "pandas"], "env_vars": {"RAY_LOG_LEVEL": "DEBUG"}},
+        metadata={"job_timeout_s": "1800", "author": "test-user"},
+        entrypoint_num_cpus=2.5,
+        entrypoint_num_gpus=1.0,
+        entrypoint_memory=2048,
+        entrypoint_resources={"custom.com/special": 1.0},
+        cluster_name="ml-cluster",
+        cluster_namespace="ml-namespace"
+    )
+    
+    job = RayJob(
+        metadata={
+            "name": "comprehensive-test-job",
+            "namespace": "test-namespace",
+            "labels": {"app": "ml-training"}
+        },
+        spec=spec
+    )
+    
+    yaml_dict = job.to_dict()
+    
+    # Verify top-level structure
+    assert yaml_dict["apiVersion"] == "ray.io/v1"
+    assert yaml_dict["kind"] == "RayJob"
+    assert yaml_dict["metadata"]["name"] == "comprehensive-test-job"
+    assert yaml_dict["metadata"]["namespace"] == "test-namespace"
+    assert yaml_dict["metadata"]["labels"] == {"app": "ml-training"}
+    
+    # Verify spec section
+    spec_dict = yaml_dict["spec"]
+    assert spec_dict["entrypoint"] == "python -c 'import ray; print(ray.cluster_resources())'"
+    assert spec_dict["submission_id"] == "test-submission-456"
+    assert spec_dict["runtime_env"] == {"pip": ["numpy==1.24.0", "pandas"], "env_vars": {"RAY_LOG_LEVEL": "DEBUG"}}
+    assert spec_dict["metadata"] == {"job_timeout_s": "1800", "author": "test-user"}
+    assert spec_dict["entrypoint_num_cpus"] == 2.5
+    assert spec_dict["entrypoint_num_gpus"] == 1.0
+    assert spec_dict["entrypoint_memory"] == 2048
+    assert spec_dict["entrypoint_resources"] == {"custom.com/special": 1.0}
+    assert spec_dict["cluster_name"] == "ml-cluster"
+    assert spec_dict["cluster_namespace"] == "ml-namespace"
+    
+    # Verify status section (should use spec status when job.status is None)
+    status_dict = yaml_dict["status"]
+    assert status_dict["status"] == RayJobStatus.PENDING
+    assert status_dict["message"] is None
+    assert status_dict["start_time"] is None
+    assert status_dict["end_time"] is None
+    assert status_dict["driver_info"] is None
+
+
+def test_ray_job_yaml_generation_minimal():
+    """Test RayJob YAML generation with minimal required fields."""
+    spec = RayJobSpec(entrypoint="python minimal_job.py")
+    
+    job = RayJob(
+        metadata={"name": "minimal-job"},
+        spec=spec
+    )
+    
+    yaml_dict = job.to_dict()
+    
+    # Verify structure
+    assert yaml_dict["apiVersion"] == "ray.io/v1"
+    assert yaml_dict["kind"] == "RayJob"
+    assert yaml_dict["metadata"]["name"] == "minimal-job"
+    
+    # Verify spec has only required field and defaults to None for others
+    spec_dict = yaml_dict["spec"]
+    assert spec_dict["entrypoint"] == "python minimal_job.py"
+    assert spec_dict["submission_id"] is None
+    assert spec_dict["runtime_env"] is None
+    assert spec_dict["metadata"] is None
+    assert spec_dict["entrypoint_num_cpus"] is None
+    assert spec_dict["entrypoint_num_gpus"] is None
+    assert spec_dict["entrypoint_memory"] is None
+    assert spec_dict["entrypoint_resources"] is None
+    assert spec_dict["cluster_name"] is None
+    assert spec_dict["cluster_namespace"] is None
+    
+    # Verify default status
+    status_dict = yaml_dict["status"]
+    assert status_dict["status"] == RayJobStatus.PENDING
+
+
+def test_ray_job_yaml_with_existing_status():
+    """Test RayJob YAML generation when status is pre-populated."""
+    spec = RayJobSpec(
+        entrypoint="python running_job.py",
+        status=RayJobStatus.PENDING  # This should be overridden by job.status
+    )
+    
+    # Simulate status from Kubernetes controller
+    existing_status = {
+        "status": "RUNNING",
+        "message": "Job is executing on cluster",
+        "start_time": "2023-12-01T10:30:00Z",
+        "end_time": None,
+        "driver_info": {
+            "id": "driver-abc123",
+            "node_ip_address": "10.244.1.5",
+            "pid": "12345"
+        }
+    }
+    
+    job = RayJob(
+        metadata={"name": "status-test-job", "namespace": "test-ns"},
+        spec=spec,
+        status=existing_status
+    )
+    
+    yaml_dict = job.to_dict()
+    
+    # Should use existing status, not spec status
+    assert yaml_dict["status"] == existing_status
+    assert yaml_dict["status"]["status"] == "RUNNING"
+    assert yaml_dict["status"]["message"] == "Job is executing on cluster"
+    assert yaml_dict["status"]["start_time"] == "2023-12-01T10:30:00Z"
+    assert yaml_dict["status"]["driver_info"]["id"] == "driver-abc123"
+
+
+def test_ray_job_yaml_with_complex_runtime_env():
+    """Test RayJob YAML generation with complex runtime environment."""
+    complex_runtime_env = {
+        "pip": ["torch==1.13.0", "transformers", "datasets"],
+        "conda": {"dependencies": ["python=3.9", "cudatoolkit=11.8"]},
+        "env_vars": {
+            "CUDA_VISIBLE_DEVICES": "0,1",
+            "PYTHONPATH": "/opt/ml/code",
+            "HF_HOME": "/tmp/huggingface"
+        },
+        "working_dir": "./training_code",
+        "py_modules": ["utils", "models"]
+    }
+    
+    spec = RayJobSpec(
+        entrypoint="python train_model.py --epochs 100",
+        runtime_env=complex_runtime_env,
+        entrypoint_num_gpus=2.0,
+        entrypoint_memory=8192
+    )
+    
+    job = RayJob(
+        metadata={"name": "complex-env-job", "namespace": "ml-training"},
+        spec=spec
+    )
+    
+    yaml_dict = job.to_dict()
+    
+    # Verify complex runtime_env is preserved exactly
+    spec_runtime_env = yaml_dict["spec"]["runtime_env"]
+    assert spec_runtime_env["pip"] == ["torch==1.13.0", "transformers", "datasets"]
+    assert spec_runtime_env["conda"]["dependencies"] == ["python=3.9", "cudatoolkit=11.8"]
+    assert spec_runtime_env["env_vars"]["CUDA_VISIBLE_DEVICES"] == "0,1"
+    assert spec_runtime_env["env_vars"]["PYTHONPATH"] == "/opt/ml/code"
+    assert spec_runtime_env["env_vars"]["HF_HOME"] == "/tmp/huggingface"
+    assert spec_runtime_env["working_dir"] == "./training_code"
+    assert spec_runtime_env["py_modules"] == ["utils", "models"]
+    
+    # Verify other fields
+    assert yaml_dict["spec"]["entrypoint"] == "python train_model.py --epochs 100"
+    assert yaml_dict["spec"]["entrypoint_num_gpus"] == 2.0
+    assert yaml_dict["spec"]["entrypoint_memory"] == 8192
+
+
+def test_ray_job_yaml_different_statuses():
+    """Test RayJob YAML generation with different status values."""
+    statuses_to_test = [
+        RayJobStatus.PENDING,
+        RayJobStatus.RUNNING, 
+        RayJobStatus.SUCCEEDED,
+        RayJobStatus.FAILED,
+        RayJobStatus.STOPPED
+    ]
+    
+    for status in statuses_to_test:
+        spec = RayJobSpec(
+            entrypoint=f"python job_{status.lower()}.py",
+            status=status
+        )
+        
+        job = RayJob(
+            metadata={"name": f"job-{status.lower()}", "namespace": "test"},
+            spec=spec
+        )
+        
+        yaml_dict = job.to_dict()
+        assert yaml_dict["status"]["status"] == status
+        assert yaml_dict["spec"]["entrypoint"] == f"python job_{status.lower()}.py"