Skip to content

Commit 0b66a1b

Browse files
committed
attempting to bump coverage
Signed-off-by: Pat O'Connor <[email protected]>
1 parent e50d3fa commit 0b66a1b

File tree

5 files changed

+339
-0
lines changed

5 files changed

+339
-0
lines changed

src/codeflare_sdk/ray/job/test_job.py

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
import pytest
2+
import yaml
3+
from pathlib import Path
4+
from unittest.mock import MagicMock, patch
5+
from codeflare_sdk.ray.job.job import RayJob, RayJobSpec, RayJobStatus
6+
7+
# Path to expected YAML files
8+
parent = Path(__file__).resolve().parents[4] # project directory
9+
expected_yamls_dir = f"{parent}/tests/test_cluster_yamls/ray"
10+
11+
12+
def test_rayjob_to_dict_minimal():
13+
"""Test RayJob.to_dict() with minimal configuration using YAML comparison"""
14+
spec = RayJobSpec(entrypoint="python test.py")
15+
job = RayJob(
16+
metadata={"name": "test-job", "namespace": "default"},
17+
spec=spec
18+
)
19+
20+
result = job.to_dict()
21+
22+
# Load expected YAML
23+
with open(f"{expected_yamls_dir}/rayjob-minimal.yaml") as f:
24+
expected = yaml.safe_load(f)
25+
26+
assert result == expected
27+
28+
29+
def test_rayjob_to_dict_full_spec():
30+
"""Test RayJob.to_dict() with complete configuration using YAML comparison"""
31+
spec = RayJobSpec(
32+
entrypoint="python main.py --config config.yaml",
33+
submission_id="job-12345",
34+
runtime_env={"working_dir": "/app", "pip": ["numpy", "pandas"]},
35+
metadata={"experiment": "test-run", "version": "1.0"},
36+
entrypoint_num_cpus=2,
37+
entrypoint_num_gpus=1,
38+
entrypoint_memory=4096,
39+
entrypoint_resources={"custom_resource": 1.0},
40+
cluster_name="test-cluster",
41+
cluster_namespace="ml-jobs",
42+
status=RayJobStatus.RUNNING,
43+
message="Job is running successfully",
44+
start_time="2024-01-01T10:00:00Z",
45+
end_time=None,
46+
driver_info={"id": "driver-123", "node_ip_address": "10.0.0.1", "pid": "12345"}
47+
)
48+
49+
job = RayJob(
50+
metadata={
51+
"name": "ml-training-job",
52+
"namespace": "ml-jobs",
53+
"labels": {"app": "ml-training", "version": "v1"},
54+
"annotations": {"description": "Machine learning training job"}
55+
},
56+
spec=spec
57+
)
58+
59+
result = job.to_dict()
60+
61+
# Load expected YAML
62+
with open(f"{expected_yamls_dir}/rayjob-full-spec.yaml") as f:
63+
expected = yaml.safe_load(f)
64+
65+
assert result == expected
66+
67+
68+
def test_rayjob_to_dict_with_existing_status():
69+
"""Test RayJob.to_dict() when status is already set using YAML comparison"""
70+
spec = RayJobSpec(entrypoint="python test.py")
71+
72+
# Pre-existing status from Kubernetes
73+
existing_status = {
74+
"jobStatus": "SUCCEEDED",
75+
"jobId": "raysubmit_12345",
76+
"message": "Job completed successfully",
77+
"startTime": "2024-01-01T10:00:00Z",
78+
"endTime": "2024-01-01T10:05:00Z"
79+
}
80+
81+
job = RayJob(
82+
metadata={"name": "completed-job", "namespace": "default"},
83+
spec=spec,
84+
status=existing_status
85+
)
86+
87+
result = job.to_dict()
88+
89+
# Load expected YAML
90+
with open(f"{expected_yamls_dir}/rayjob-existing-status.yaml") as f:
91+
expected = yaml.safe_load(f)
92+
93+
assert result == expected
94+
95+
96+
def test_rayjob_status_enum_values():
97+
"""Test all RayJobStatus enum values"""
98+
assert RayJobStatus.PENDING == "PENDING"
99+
assert RayJobStatus.RUNNING == "RUNNING"
100+
assert RayJobStatus.STOPPED == "STOPPED"
101+
assert RayJobStatus.SUCCEEDED == "SUCCEEDED"
102+
assert RayJobStatus.FAILED == "FAILED"
103+
104+
105+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
106+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
107+
@patch('kubernetes.dynamic.DynamicClient')
108+
def test_rayjob_apply_success(mock_dynamic_client, mock_get_api_client, mock_config_check):
109+
"""Test RayJob.apply() method successful execution"""
110+
# Mock the Kubernetes API components
111+
mock_api_instance = MagicMock()
112+
mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
113+
114+
spec = RayJobSpec(entrypoint="python test.py")
115+
job = RayJob(
116+
metadata={"name": "test-job", "namespace": "test-ns"},
117+
spec=spec
118+
)
119+
120+
# Test successful apply
121+
job.apply()
122+
123+
# Verify the API was called correctly
124+
mock_config_check.assert_called_once()
125+
mock_get_api_client.assert_called_once()
126+
mock_dynamic_client.assert_called_once()
127+
mock_api_instance.server_side_apply.assert_called_once()
128+
129+
# Check the server_side_apply call arguments
130+
call_args = mock_api_instance.server_side_apply.call_args
131+
assert call_args[1]['field_manager'] == 'codeflare-sdk'
132+
assert call_args[1]['group'] == 'ray.io'
133+
assert call_args[1]['version'] == 'v1'
134+
assert call_args[1]['namespace'] == 'test-ns'
135+
assert call_args[1]['plural'] == 'rayjobs'
136+
assert call_args[1]['force_conflicts'] == False
137+
138+
# Verify the body contains the expected job structure
139+
body = call_args[1]['body']
140+
assert body['apiVersion'] == 'ray.io/v1'
141+
assert body['kind'] == 'RayJob'
142+
assert body['metadata']['name'] == 'test-job'
143+
assert body['spec']['entrypoint'] == 'python test.py'
144+
145+
146+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
147+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
148+
@patch('kubernetes.dynamic.DynamicClient')
149+
def test_rayjob_apply_with_force(mock_dynamic_client, mock_get_api_client, mock_config_check):
150+
"""Test RayJob.apply() method with force=True"""
151+
mock_api_instance = MagicMock()
152+
mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
153+
154+
spec = RayJobSpec(entrypoint="python test.py")
155+
job = RayJob(
156+
metadata={"name": "test-job", "namespace": "default"},
157+
spec=spec
158+
)
159+
160+
# Test apply with force=True
161+
job.apply(force=True)
162+
163+
# Verify force_conflicts was set to True
164+
call_args = mock_api_instance.server_side_apply.call_args
165+
assert call_args[1]['force_conflicts'] == True
166+
167+
168+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
169+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
170+
@patch('kubernetes.dynamic.DynamicClient')
171+
def test_rayjob_apply_dynamic_client_error(mock_dynamic_client, mock_get_api_client, mock_config_check):
172+
"""Test RayJob.apply() method with DynamicClient initialization error"""
173+
# Mock DynamicClient to raise AttributeError
174+
mock_dynamic_client.side_effect = AttributeError("Failed to connect")
175+
176+
spec = RayJobSpec(entrypoint="python test.py")
177+
job = RayJob(
178+
metadata={"name": "test-job", "namespace": "default"},
179+
spec=spec
180+
)
181+
182+
# Test that RuntimeError is raised
183+
with pytest.raises(RuntimeError, match="Failed to initialize DynamicClient"):
184+
job.apply()
185+
186+
187+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check')
188+
@patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client')
189+
@patch('kubernetes.dynamic.DynamicClient')
190+
@patch('codeflare_sdk.common._kube_api_error_handling')
191+
def test_rayjob_apply_api_error(mock_error_handling, mock_dynamic_client, mock_get_api_client, mock_config_check):
192+
"""Test RayJob.apply() method with Kubernetes API error"""
193+
# Mock the API to raise an exception
194+
mock_api_instance = MagicMock()
195+
mock_api_instance.server_side_apply.side_effect = Exception("API Error")
196+
mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
197+
198+
spec = RayJobSpec(entrypoint="python test.py")
199+
job = RayJob(
200+
metadata={"name": "test-job", "namespace": "default"},
201+
spec=spec
202+
)
203+
204+
# Test that error handling is called
205+
job.apply()
206+
207+
# Verify error handling was called
208+
mock_error_handling.assert_called_once()
209+
210+
211+
def test_rayjob_default_namespace_in_apply():
212+
"""Test that RayJob.apply() uses 'default' namespace when not specified in metadata"""
213+
with patch('codeflare_sdk.common.kubernetes_cluster.auth.config_check'), \
214+
patch('codeflare_sdk.common.kubernetes_cluster.auth.get_api_client'), \
215+
patch('kubernetes.dynamic.DynamicClient') as mock_dynamic_client:
216+
217+
mock_api_instance = MagicMock()
218+
mock_dynamic_client.return_value.resources.get.return_value = mock_api_instance
219+
220+
spec = RayJobSpec(entrypoint="python test.py")
221+
job = RayJob(
222+
metadata={"name": "test-job"}, # No namespace specified
223+
spec=spec
224+
)
225+
226+
job.apply()
227+
228+
# Verify default namespace was used
229+
call_args = mock_api_instance.server_side_apply.call_args
230+
assert call_args[1]['namespace'] == 'default'

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
import pytest
1516
from unittest.mock import MagicMock
1617
from codeflare_sdk.ray.rayjobs.rayjob import RayJob
1718

@@ -58,3 +59,30 @@ def test_rayjob_submit_success(mocker):
5859
assert job_cr["spec"]["entrypoint"] == "python -c 'print(\"hello world\")'"
5960
assert job_cr["spec"]["clusterSelector"]["ray.io/cluster"] == "test-ray-cluster"
6061
assert job_cr["spec"]["runtimeEnvYAML"] == "{'pip': ['requests']}"
62+
63+
64+
def test_rayjob_submit_failure(mocker):
65+
"""Test RayJob submission failure."""
66+
# Mock kubernetes config loading
67+
mocker.patch("kubernetes.config.load_kube_config")
68+
69+
# Mock the RayjobApi class entirely
70+
mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi")
71+
mock_api_instance = MagicMock()
72+
mock_api_class.return_value = mock_api_instance
73+
74+
# Configure the mock to return failure (False/None) when submit_job is called
75+
mock_api_instance.submit_job.return_value = None
76+
77+
# Create a RayJob instance
78+
rayjob = RayJob(
79+
job_name="test-rayjob",
80+
cluster_name="test-ray-cluster",
81+
namespace="default",
82+
entrypoint="python script.py",
83+
runtime_env={"pip": ["numpy"]},
84+
)
85+
86+
# Test that RuntimeError is raised on failure
87+
with pytest.raises(RuntimeError, match="Failed to submit RayJob test-rayjob"):
88+
rayjob.submit()
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
apiVersion: ray.io/v1
2+
kind: RayJob
3+
metadata:
4+
name: completed-job
5+
namespace: default
6+
spec:
7+
entrypoint: python test.py
8+
submission_id: null
9+
runtime_env: null
10+
metadata: null
11+
entrypoint_num_cpus: null
12+
entrypoint_num_gpus: null
13+
entrypoint_memory: null
14+
entrypoint_resources: null
15+
cluster_name: null
16+
cluster_namespace: null
17+
status:
18+
jobStatus: SUCCEEDED
19+
jobId: raysubmit_12345
20+
message: Job completed successfully
21+
startTime: '2024-01-01T10:00:00Z'
22+
endTime: '2024-01-01T10:05:00Z'
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
apiVersion: ray.io/v1
2+
kind: RayJob
3+
metadata:
4+
name: ml-training-job
5+
namespace: ml-jobs
6+
labels:
7+
app: ml-training
8+
version: v1
9+
annotations:
10+
description: Machine learning training job
11+
spec:
12+
entrypoint: python main.py --config config.yaml
13+
submission_id: job-12345
14+
runtime_env:
15+
working_dir: /app
16+
pip:
17+
- numpy
18+
- pandas
19+
metadata:
20+
experiment: test-run
21+
version: '1.0'
22+
entrypoint_num_cpus: 2
23+
entrypoint_num_gpus: 1
24+
entrypoint_memory: 4096
25+
entrypoint_resources:
26+
custom_resource: 1.0
27+
cluster_name: test-cluster
28+
cluster_namespace: ml-jobs
29+
status:
30+
status: RUNNING
31+
message: Job is running successfully
32+
start_time: '2024-01-01T10:00:00Z'
33+
end_time: null
34+
driver_info:
35+
id: driver-123
36+
node_ip_address: 10.0.0.1
37+
pid: '12345'
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
apiVersion: ray.io/v1
2+
kind: RayJob
3+
metadata:
4+
name: test-job
5+
namespace: default
6+
spec:
7+
entrypoint: python test.py
8+
submission_id: null
9+
runtime_env: null
10+
metadata: null
11+
entrypoint_num_cpus: null
12+
entrypoint_num_gpus: null
13+
entrypoint_memory: null
14+
entrypoint_resources: null
15+
cluster_name: null
16+
cluster_namespace: null
17+
status:
18+
status: PENDING
19+
message: null
20+
start_time: null
21+
end_time: null
22+
driver_info: null

0 commit comments

Comments
 (0)