From 3c37651345d45db95ab4c272e8eb0658447fc818 Mon Sep 17 00:00:00 2001
From: MichaelClifford <mcliffor@redhat.com>
Date: Fri, 16 Jun 2023 17:17:39 -0400
Subject: [PATCH] make priority levels global

---
 src/codeflare_sdk/utils/generate_yaml.py | 17 +++++++++--------
 tests/unit_test.py                       |  1 -
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/codeflare_sdk/utils/generate_yaml.py b/src/codeflare_sdk/utils/generate_yaml.py
index e5389cca4..5e2e7cb4e 100755
--- a/src/codeflare_sdk/utils/generate_yaml.py
+++ b/src/codeflare_sdk/utils/generate_yaml.py
@@ -76,17 +76,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
index 643c408d1..a743670f2 100644
--- a/tests/unit_test.py
+++ b/tests/unit_test.py
@@ -234,7 +234,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(