Skip to content

Commit 2aa287b

Browse files
committed
Updated unit test for user labels
1 parent 7991146 commit 2aa287b

File tree

3 files changed

+12
-7
lines changed

3 files changed

+12
-7
lines changed

Diff for: src/codeflare_sdk/utils/generate_yaml.py

+7-7
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,7 @@ def write_components(
313313
output_file_name: str,
314314
namespace: str,
315315
local_queue: Optional[str],
316-
user_labels: dict,
316+
cluster_labels: dict,
317317
):
318318
# Create the directory if it doesn't exist
319319
directory_path = os.path.dirname(output_file_name)
@@ -335,7 +335,7 @@ def write_components(
335335
]
336336
labels = component["generictemplate"]["metadata"]["labels"]
337337
labels.update({"kueue.x-k8s.io/queue-name": lq_name})
338-
labels.update(user_labels)
338+
labels.update(cluster_labels)
339339
outfile.write("---\n")
340340
yaml.dump(
341341
component["generictemplate"], outfile, default_flow_style=False
@@ -348,7 +348,7 @@ def load_components(
348348
name: str,
349349
namespace: str,
350350
local_queue: Optional[str],
351-
user_labels: dict,
351+
cluster_labels: dict,
352352
):
353353
component_list = []
354354
components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
@@ -364,7 +364,7 @@ def load_components(
364364
]
365365
labels = component["generictemplate"]["metadata"]["labels"]
366366
labels.update({"kueue.x-k8s.io/queue-name": lq_name})
367-
labels.update(user_labels)
367+
labels.update(cluster_labels)
368368
component_list.append(component["generictemplate"])
369369

370370
resources = "---\n" + "---\n".join(
@@ -405,7 +405,7 @@ def generate_appwrapper(
405405
write_to_file: bool,
406406
verify_tls: bool,
407407
local_queue: Optional[str],
408-
user_labels,
408+
cluster_labels,
409409
):
410410
user_yaml = read_template(template)
411411
appwrapper_name, cluster_name = gen_names(name)
@@ -457,13 +457,13 @@ def generate_appwrapper(
457457
if mcad:
458458
write_user_appwrapper(user_yaml, outfile)
459459
else:
460-
write_components(user_yaml, outfile, namespace, local_queue, user_labels)
460+
write_components(user_yaml, outfile, namespace, local_queue, cluster_labels)
461461
return outfile
462462
else:
463463
if mcad:
464464
user_yaml = load_appwrapper(user_yaml, name)
465465
else:
466466
user_yaml = load_components(
467-
user_yaml, name, namespace, local_queue, user_labels
467+
user_yaml, name, namespace, local_queue, cluster_labels
468468
)
469469
return user_yaml

Diff for: tests/test-case-no-mcad.yamls

+2
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ metadata:
55
labels:
66
controller-tools.k8s.io: '1.0'
77
kueue.x-k8s.io/queue-name: local-queue-default
8+
testlabel: test
9+
testlabel2: test
810
name: unit-test-cluster-ray
911
namespace: ns
1012
spec:

Diff for: tests/unit_test.py

+3
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,7 @@ def test_cluster_creation_no_mcad(mocker):
324324
config.name = "unit-test-cluster-ray"
325325
config.write_to_file = True
326326
config.mcad = False
327+
config.cluster_labels = {"testlabel": "test", "testlabel2": "test"}
327328
cluster = Cluster(config)
328329

329330
assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
@@ -348,6 +349,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
348349
config.mcad = False
349350
config.write_to_file = True
350351
config.local_queue = "local-queue-default"
352+
config.cluster_labels = {"testlabel": "test", "testlabel2": "test"}
351353
cluster = Cluster(config)
352354
assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
353355
assert cluster.app_wrapper_name == "unit-test-cluster-ray"
@@ -373,6 +375,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
373375
write_to_file=True,
374376
mcad=False,
375377
local_queue="local-queue-default",
378+
cluster_labels={"testlabel": "test", "testlabel2": "test"},
376379
)
377380
cluster = Cluster(config)
378381
assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"

0 commit comments

Comments
 (0)