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'
0 commit comments