make priority levels global

project-codeflare · Jul 26, 2023 · a1f084c · a1f084c
1 parent 1aefc04
commit a1f084c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/src/codeflare_sdk/utils/generate_yaml.py b/src/codeflare_sdk/utils/generate_yaml.py
@@ -88,17 +88,18 @@ def update_labels(yaml, instascale, instance_types):
         metadata.pop("labels")
 
 
+PRIORITY_LEVELS = {
+    "low": (1, "low-priority"),
+    "default": (5, "default-priority"),
+    "high": (10, "high-priority"),
+}
+
+
 def update_priority(yaml, item, priority):
-    if priority not in ["low", "default", "high"]:
+    if priority not in PRIORITY_LEVELS:
         sys.exit("Priority must be 'low', 'default', or 'high'")
 
-    priority_levels = {
-        "low": (1, "low-priority"),
-        "default": (5, "default-priority"),
-        "high": (10, "high-priority"),
-    }
-
-    priority_level = priority_levels[priority]
+    priority_level = PRIORITY_LEVELS[priority]
     spec = yaml.get("spec")
     spec["priority"] = priority_level[0]
     # spec["SchedulingSpec"]["priorityClassName"] = priority_level

diff --git a/tests/unit_test.py b/tests/unit_test.py
@@ -202,7 +202,6 @@ def test_config_creation():
 
 def test_cluster_creation():
     cluster = Cluster(test_config_creation())
-    print(cluster.app_wrapper_yaml)
     assert cluster.app_wrapper_yaml == "unit-test-cluster.yaml"
     assert cluster.app_wrapper_name == "unit-test-cluster"
     assert filecmp.cmp(