11
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
- from ..utils .unit_test_support import (
15
- apply_template ,
16
- get_local_queue ,
17
- createClusterConfig ,
18
- get_template_variables ,
19
- )
14
+ from ..utils .unit_test_support import get_local_queue , create_cluster_config , get_template_variables , apply_template
20
15
from unittest .mock import patch
21
16
from codeflare_sdk .ray .cluster .cluster import Cluster , ClusterConfiguration
22
17
import yaml
23
18
import os
24
19
import filecmp
25
20
from pathlib import Path
26
- from .kueue import list_local_queues
21
+ from .kueue import list_local_queues , local_queue_exists , add_queue_label
27
22
28
23
parent = Path (__file__ ).resolve ().parents [4 ] # project directory
29
24
aw_dir = os .path .expanduser ("~/.codeflare/resources/" )
@@ -51,7 +46,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
51
46
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object" ,
52
47
return_value = get_local_queue ("kueue.x-k8s.io" , "v1beta1" , "ns" , "localqueues" ),
53
48
)
54
- config = createClusterConfig ()
49
+ config = create_cluster_config ()
55
50
config .name = "unit-test-cluster-kueue"
56
51
config .write_to_file = True
57
52
config .local_queue = "local-queue-default"
@@ -67,7 +62,7 @@ def test_cluster_creation_no_aw_local_queue(mocker):
67
62
assert cluster_kueue == expected_rc
68
63
69
64
# With resources loaded in memory, no Local Queue specified.
70
- config = createClusterConfig ()
65
+ config = create_cluster_config ()
71
66
config .name = "unit-test-cluster-kueue"
72
67
config .write_to_file = False
73
68
cluster = Cluster (config )
@@ -84,7 +79,7 @@ def test_aw_creation_local_queue(mocker):
84
79
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object" ,
85
80
return_value = get_local_queue ("kueue.x-k8s.io" , "v1beta1" , "ns" , "localqueues" ),
86
81
)
87
- config = createClusterConfig ()
82
+ config = create_cluster_config ()
88
83
config .name = "unit-test-aw-kueue"
89
84
config .appwrapper = True
90
85
config .write_to_file = True
@@ -101,7 +96,7 @@ def test_aw_creation_local_queue(mocker):
101
96
assert aw_kueue == expected_rc
102
97
103
98
# With resources loaded in memory, no Local Queue specified.
104
- config = createClusterConfig ()
99
+ config = create_cluster_config ()
105
100
config .name = "unit-test-aw-kueue"
106
101
config .appwrapper = True
107
102
config .write_to_file = False
@@ -120,7 +115,7 @@ def test_get_local_queue_exists_fail(mocker):
120
115
"kubernetes.client.CustomObjectsApi.list_namespaced_custom_object" ,
121
116
return_value = get_local_queue ("kueue.x-k8s.io" , "v1beta1" , "ns" , "localqueues" ),
122
117
)
123
- config = createClusterConfig ()
118
+ config = create_cluster_config ()
124
119
config .name = "unit-test-aw-kueue"
125
120
config .appwrapper = True
126
121
config .write_to_file = True
@@ -175,6 +170,123 @@ def test_list_local_queues(mocker):
175
170
assert lqs == []
176
171
177
172
173
+ def test_local_queue_exists_found (mocker ):
174
+ # Mock Kubernetes client and list_namespaced_custom_object method
175
+ mocker .patch ("kubernetes.config.load_kube_config" , return_value = "ignore" )
176
+ mock_api_instance = mocker .Mock ()
177
+ mocker .patch ("kubernetes.client.CustomObjectsApi" , return_value = mock_api_instance )
178
+ mocker .patch ("codeflare_sdk.ray.cluster.cluster.config_check" )
179
+
180
+ # Mock return value for list_namespaced_custom_object
181
+ mock_api_instance .list_namespaced_custom_object .return_value = {
182
+ "items" : [
183
+ {"metadata" : {"name" : "existing-queue" }},
184
+ {"metadata" : {"name" : "another-queue" }},
185
+ ]
186
+ }
187
+
188
+ # Call the function
189
+ namespace = "test-namespace"
190
+ local_queue_name = "existing-queue"
191
+ result = local_queue_exists (namespace , local_queue_name )
192
+
193
+ # Assertions
194
+ assert result is True
195
+ mock_api_instance .list_namespaced_custom_object .assert_called_once_with (
196
+ group = "kueue.x-k8s.io" ,
197
+ version = "v1beta1" ,
198
+ namespace = namespace ,
199
+ plural = "localqueues" ,
200
+ )
201
+
202
+
203
+ def test_local_queue_exists_not_found (mocker ):
204
+ # Mock Kubernetes client and list_namespaced_custom_object method
205
+ mocker .patch ("kubernetes.config.load_kube_config" , return_value = "ignore" )
206
+ mock_api_instance = mocker .Mock ()
207
+ mocker .patch ("kubernetes.client.CustomObjectsApi" , return_value = mock_api_instance )
208
+ mocker .patch ("codeflare_sdk.ray.cluster.cluster.config_check" )
209
+
210
+ # Mock return value for list_namespaced_custom_object
211
+ mock_api_instance .list_namespaced_custom_object .return_value = {
212
+ "items" : [
213
+ {"metadata" : {"name" : "another-queue" }},
214
+ {"metadata" : {"name" : "different-queue" }},
215
+ ]
216
+ }
217
+
218
+ # Call the function
219
+ namespace = "test-namespace"
220
+ local_queue_name = "non-existent-queue"
221
+ result = local_queue_exists (namespace , local_queue_name )
222
+
223
+ # Assertions
224
+ assert result is False
225
+ mock_api_instance .list_namespaced_custom_object .assert_called_once_with (
226
+ group = "kueue.x-k8s.io" ,
227
+ version = "v1beta1" ,
228
+ namespace = namespace ,
229
+ plural = "localqueues" ,
230
+ )
231
+
232
+
233
+ import pytest
234
+ from unittest import mock # If you're also using mocker from pytest-mock
235
+
236
+
237
+ def test_add_queue_label_with_valid_local_queue (mocker ):
238
+ # Mock the kubernetes.client.CustomObjectsApi and its response
239
+ mock_api_instance = mocker .patch ("kubernetes.client.CustomObjectsApi" )
240
+ mock_api_instance .return_value .list_namespaced_custom_object .return_value = {
241
+ "items" : [
242
+ {"metadata" : {"name" : "valid-queue" }},
243
+ ]
244
+ }
245
+
246
+ # Mock other dependencies
247
+ mocker .patch ("codeflare_sdk.common.kueue.local_queue_exists" , return_value = True )
248
+ mocker .patch (
249
+ "codeflare_sdk.common.kueue.get_default_kueue_name" ,
250
+ return_value = "default-queue" ,
251
+ )
252
+
253
+ # Define input item and parameters
254
+ item = {"metadata" : {}}
255
+ namespace = "test-namespace"
256
+ local_queue = "valid-queue"
257
+
258
+ # Call the function
259
+ add_queue_label (item , namespace , local_queue )
260
+
261
+ # Assert that the label is added to the item
262
+ assert item ["metadata" ]["labels" ] == {"kueue.x-k8s.io/queue-name" : "valid-queue" }
263
+
264
+
265
+ def test_add_queue_label_with_invalid_local_queue (mocker ):
266
+ # Mock the kubernetes.client.CustomObjectsApi and its response
267
+ mock_api_instance = mocker .patch ("kubernetes.client.CustomObjectsApi" )
268
+ mock_api_instance .return_value .list_namespaced_custom_object .return_value = {
269
+ "items" : [
270
+ {"metadata" : {"name" : "valid-queue" }},
271
+ ]
272
+ }
273
+
274
+ # Mock the local_queue_exists function to return False
275
+ mocker .patch ("codeflare_sdk.common.kueue.local_queue_exists" , return_value = False )
276
+
277
+ # Define input item and parameters
278
+ item = {"metadata" : {}}
279
+ namespace = "test-namespace"
280
+ local_queue = "invalid-queue"
281
+
282
+ # Call the function and expect a ValueError
283
+ with pytest .raises (
284
+ ValueError ,
285
+ match = "local_queue provided does not exist or is not in this namespace" ,
286
+ ):
287
+ add_queue_label (item , namespace , local_queue )
288
+
289
+
178
290
# Make sure to always keep this function last
179
291
def test_cleanup ():
180
292
os .remove (f"{ aw_dir } unit-test-cluster-kueue.yaml" )
0 commit comments