test refactor

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

Author: chipspeak

src/codeflare_sdk/ray/job/test_job.py

@@ -1,28 +1,38 @@
+# Copyright 2024 IBM, Red Hat
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import pytest
 import yaml
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 from codeflare_sdk.ray.job.job import RayJob, RayJobSpec, RayJobStatus
 
-# Path to expected YAML files
 parent = Path(__file__).resolve().parents[4]  # project directory
 expected_yamls_dir = f"{parent}/tests/test_cluster_yamls/ray"
 
 
 def test_rayjob_to_dict_minimal():
     """Test RayJob.to_dict() with minimal configuration using YAML comparison"""
     spec = RayJobSpec(entrypoint="python test.py")
-    job = RayJob(
-        metadata={"name": "test-job", "namespace": "default"},
-        spec=spec
-    )
-    
+    job = RayJob(metadata={"name": "test-job", "namespace": "default"}, spec=spec)
+
     result = job.to_dict()
-    
+
     # Load expected YAML
     with open(f"{expected_yamls_dir}/rayjob-minimal.yaml") as f:
         expected = yaml.safe_load(f)
-    
+
     assert result == expected
 
 
@@ -43,53 +53,53 @@ def test_rayjob_to_dict_full_spec():
         message="Job is running successfully",
         start_time="2024-01-01T10:00:00Z",
         end_time=None,
-        driver_info={"id": "driver-123", "node_ip_address": "10.0.0.1", "pid": "12345"}
+        driver_info={"id": "driver-123", "node_ip_address": "10.0.0.1", "pid": "12345"},
     )
-    
+
     job = RayJob(
         metadata={
             "name": "ml-training-job",
             "namespace": "ml-jobs",
             "labels": {"app": "ml-training", "version": "v1"},
-            "annotations": {"description": "Machine learning training job"}
+            "annotations": {"description": "Machine learning training job"},
         },
-        spec=spec
+        spec=spec,
     )
-    
+
     result = job.to_dict()
-    
+
     # Load expected YAML
     with open(f"{expected_yamls_dir}/rayjob-full-spec.yaml") as f:
         expected = yaml.safe_load(f)
-    
+
     assert result == expected
 
 
 def test_rayjob_to_dict_with_existing_status():
     """Test RayJob.to_dict() when status is already set using YAML comparison"""
     spec = RayJobSpec(entrypoint="python test.py")
-    
+
     # Pre-existing status from Kubernetes
     existing_status = {
         "jobStatus": "SUCCEEDED",
         "jobId": "raysubmit_12345",
         "message": "Job completed successfully",
         "startTime": "2024-01-01T10:00:00Z",
-        "endTime": "2024-01-01T10:05:00Z"
+        "endTime": "2024-01-01T10:05:00Z",
     }
-    
+
     job = RayJob(
         metadata={"name": "completed-job", "namespace": "default"},
         spec=spec,
-        status=existing_status
+        status=existing_status,
     )
-    
+
     result = job.to_dict()
-    
+
     # Load expected YAML
     with open(f"{expected_yamls_dir}/rayjob-existing-status.yaml") as f:
         expected = yaml.safe_load(f)
-    
+
     assert result == expected
 
 
@@ -102,129 +112,119 @@ def test_rayjob_status_enum_values():
     assert RayJobStatus.FAILED == "FAILED"
 
 
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
-@patch('kubernetes.dynamic.DynamicClient')
-def test_rayjob_apply_success(mock_dynamic_client, mock_get_api_client, mock_config_check):
+def test_rayjob_apply_success(mocker):
     """Test RayJob.apply() method successful execution"""
-    # Mock the Kubernetes API components
+    # Mock Kubernetes configuration and API components
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_check")
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client")
+    mock_dynamic_client = mocker.patch("kubernetes.dynamic.DynamicClient")
+
     mock_api_instance = MagicMock()
     mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
-    
+
     spec = RayJobSpec(entrypoint="python test.py")
-    job = RayJob(
-        metadata={"name": "test-job", "namespace": "test-ns"},
-        spec=spec
-    )
-    
+    job = RayJob(metadata={"name": "test-job", "namespace": "test-ns"}, spec=spec)
+
     # Test successful apply
     job.apply()
-    
+
     # Verify the API was called correctly
-    mock_config_check.assert_called_once()
-    mock_get_api_client.assert_called_once()
-    mock_dynamic_client.assert_called_once()
     mock_api_instance.server_side_apply.assert_called_once()
-    
+
     # Check the server_side_apply call arguments
     call_args = mock_api_instance.server_side_apply.call_args
-    assert call_args[1]['field_manager'] == 'codeflare-sdk'
-    assert call_args[1]['group'] == 'ray.io'
-    assert call_args[1]['version'] == 'v1'
-    assert call_args[1]['namespace'] == 'test-ns'
-    assert call_args[1]['plural'] == 'rayjobs'
-    assert call_args[1]['force_conflicts'] == False
-    
+    assert call_args[1]["field_manager"] == "codeflare-sdk"
+    assert call_args[1]["group"] == "ray.io"
+    assert call_args[1]["version"] == "v1"
+    assert call_args[1]["namespace"] == "test-ns"
+    assert call_args[1]["plural"] == "rayjobs"
+    assert call_args[1]["force_conflicts"] == False
+
     # Verify the body contains the expected job structure
-    body = call_args[1]['body']
-    assert body['apiVersion'] == 'ray.io/v1'
-    assert body['kind'] == 'RayJob'
-    assert body['metadata']['name'] == 'test-job'
-    assert body['spec']['entrypoint'] == 'python test.py'
+    body = call_args[1]["body"]
+    assert body["apiVersion"] == "ray.io/v1"
+    assert body["kind"] == "RayJob"
+    assert body["metadata"]["name"] == "test-job"
+    assert body["spec"]["entrypoint"] == "python test.py"
 
 
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
-@patch('kubernetes.dynamic.DynamicClient')
-def test_rayjob_apply_with_force(mock_dynamic_client, mock_get_api_client, mock_config_check):
+def test_rayjob_apply_with_force(mocker):
     """Test RayJob.apply() method with force=True"""
+    # Mock Kubernetes configuration and API components
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_check")
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client")
+    mock_dynamic_client = mocker.patch("kubernetes.dynamic.DynamicClient")
+
     mock_api_instance = MagicMock()
     mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
-    
+
     spec = RayJobSpec(entrypoint="python test.py")
-    job = RayJob(
-        metadata={"name": "test-job", "namespace": "default"},
-        spec=spec
-    )
-    
+    job = RayJob(metadata={"name": "test-job", "namespace": "default"}, spec=spec)
+
     # Test apply with force=True
     job.apply(force=True)
-    
+
     # Verify force_conflicts was set to True
     call_args = mock_api_instance.server_side_apply.call_args
-    assert call_args[1]['force_conflicts'] == True
+    assert call_args[1]["force_conflicts"] == True
 
 
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
-@patch('kubernetes.dynamic.DynamicClient')
-def test_rayjob_apply_dynamic_client_error(mock_dynamic_client, mock_get_api_client, mock_config_check):
+def test_rayjob_apply_dynamic_client_error(mocker):
     """Test RayJob.apply() method with DynamicClient initialization error"""
+    # Mock Kubernetes configuration
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_check")
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client")
+
     # Mock DynamicClient to raise AttributeError
+    mock_dynamic_client = mocker.patch("kubernetes.dynamic.DynamicClient")
     mock_dynamic_client.side_effect = AttributeError("Failed to connect")
-    
+
     spec = RayJobSpec(entrypoint="python test.py")
-    job = RayJob(
-        metadata={"name": "test-job", "namespace": "default"},
-        spec=spec
-    )
-    
+    job = RayJob(metadata={"name": "test-job", "namespace": "default"}, spec=spec)
+
     # Test that RuntimeError is raised
     with pytest.raises(RuntimeError, match="Failed to initialize DynamicClient"):
         job.apply()
 
 
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
-@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
-@patch('kubernetes.dynamic.DynamicClient')
-@patch('codeflare_sdk.common._kube_api_error_handling')
-def test_rayjob_apply_api_error(mock_error_handling, mock_dynamic_client, mock_get_api_client, mock_config_check):
+def test_rayjob_apply_api_error(mocker):
     """Test RayJob.apply() method with Kubernetes API error"""
+    # Mock Kubernetes configuration and API components
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_check")
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client")
+    mock_dynamic_client = mocker.patch("kubernetes.dynamic.DynamicClient")
+    mock_error_handling = mocker.patch("codeflare_sdk.common._kube_api_error_handling")
+
     # Mock the API to raise an exception
     mock_api_instance = MagicMock()
     mock_api_instance.server_side_apply.side_effect = Exception("API Error")
     mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
-    
+
     spec = RayJobSpec(entrypoint="python test.py")
-    job = RayJob(
-        metadata={"name": "test-job", "namespace": "default"},
-        spec=spec
-    )
-    
+    job = RayJob(metadata={"name": "test-job", "namespace": "default"}, spec=spec)
+
     # Test that error handling is called
     job.apply()
-    
+
     # Verify error handling was called
     mock_error_handling.assert_called_once()
 
 
-def test_rayjob_default_namespace_in_apply():
+def test_rayjob_default_namespace_in_apply(mocker):
     """Test that RayJob.apply() uses 'default' namespace when not specified in metadata"""
-    with patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check'), \
-         patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client'), \
-         patch('kubernetes.dynamic.DynamicClient') as mock_dynamic_client:
-        
-        mock_api_instance = MagicMock()
-        mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
-        
-        spec = RayJobSpec(entrypoint="python test.py")
-        job = RayJob(
-            metadata={"name": "test-job"},  # No namespace specified
-            spec=spec
-        )
-        
-        job.apply()
-        
-        # Verify default namespace was used
-        call_args = mock_api_instance.server_side_apply.call_args
-        assert call_args[1]['namespace'] == 'default' 
+    # Mock Kubernetes configuration and API components
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.config_check")
+    mocker.patch("codeflare_sdk.common.kubernetes_cluster.auth.get_api_client")
+    mock_dynamic_client = mocker.patch("kubernetes.dynamic.DynamicClient")
+
+    mock_api_instance = MagicMock()
+    mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
+
+    spec = RayJobSpec(entrypoint="python test.py")
+    job = RayJob(metadata={"name": "test-job"}, spec=spec)  # No namespace specified
+
+    job.apply()
+
+    # Verify default namespace was used
+    call_args = mock_api_instance.server_side_apply.call_args
+    assert call_args[1]["namespace"] == "default"

tests/test_cluster_yamls/ray/rayjob-existing-status.yaml

@@ -19,4 +19,4 @@ status:
   jobId: raysubmit_12345
   message: Job completed successfully
   startTime: '2024-01-01T10:00:00Z'
-  endTime: '2024-01-01T10:05:00Z' 
+  endTime: '2024-01-01T10:05:00Z'

tests/test_cluster_yamls/ray/rayjob-full-spec.yaml

@@ -34,4 +34,4 @@ status:
   driver_info:
     id: driver-123
     node_ip_address: 10.0.0.1
-    pid: '12345' 
+    pid: '12345'

tests/test_cluster_yamls/ray/rayjob-minimal.yaml

@@ -19,4 +19,4 @@ status:
   message: null
   start_time: null
   end_time: null
-  driver_info: null 
+  driver_info: null