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