update unit tests to work on both Python 3.9 and 3.11

Author: Ygnas

src/codeflare_sdk/common/kueue/test_kueue.py

@@ -11,7 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from ..utils.unit_test_support import get_local_queue, createClusterConfig
+from ..utils.unit_test_support import (
+    apply_template,
+    get_local_queue,
+    createClusterConfig,
+    get_template_variables,
+)
 from unittest.mock import patch
 from codeflare_sdk.ray.cluster.cluster import Cluster, ClusterConfiguration
 import yaml
@@ -52,21 +57,21 @@ def test_cluster_creation_no_aw_local_queue(mocker):
     config.local_queue = "local-queue-default"
     cluster = Cluster(config)
     assert cluster.resource_yaml == f"{aw_dir}unit-test-cluster-kueue.yaml"
-    assert filecmp.cmp(
-        f"{aw_dir}unit-test-cluster-kueue.yaml",
+    expected_rc = apply_template(
         f"{parent}/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml",
-        shallow=True,
+        get_template_variables(),
     )
 
+    with open(f"{aw_dir}unit-test-cluster-kueue.yaml", "r") as f:
+        cluster_kueue = yaml.load(f, Loader=yaml.FullLoader)
+        assert cluster_kueue == expected_rc
+
     # With resources loaded in memory, no Local Queue specified.
     config = createClusterConfig()
     config.name = "unit-test-cluster-kueue"
     config.write_to_file = False
     cluster = Cluster(config)
-
-    with open(f"{parent}/tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml") as f:
-        expected_rc = yaml.load(f, Loader=yaml.FullLoader)
-        assert cluster.resource_yaml == expected_rc
+    assert cluster.resource_yaml == expected_rc
 
 
 def test_aw_creation_local_queue(mocker):
@@ -86,22 +91,23 @@ def test_aw_creation_local_queue(mocker):
     config.local_queue = "local-queue-default"
     cluster = Cluster(config)
     assert cluster.resource_yaml == f"{aw_dir}unit-test-aw-kueue.yaml"
-    assert filecmp.cmp(
-        f"{aw_dir}unit-test-aw-kueue.yaml",
+    expected_rc = apply_template(
         f"{parent}/tests/test_cluster_yamls/kueue/aw_kueue.yaml",
-        shallow=True,
+        get_template_variables(),
     )
 
+    with open(f"{aw_dir}unit-test-aw-kueue.yaml", "r") as f:
+        aw_kueue = yaml.load(f, Loader=yaml.FullLoader)
+        assert aw_kueue == expected_rc
+
     # With resources loaded in memory, no Local Queue specified.
     config = createClusterConfig()
     config.name = "unit-test-aw-kueue"
     config.appwrapper = True
     config.write_to_file = False
     cluster = Cluster(config)
 
-    with open(f"{parent}/tests/test_cluster_yamls/kueue/aw_kueue.yaml") as f:
-        expected_rc = yaml.load(f, Loader=yaml.FullLoader)
-        assert cluster.resource_yaml == expected_rc
+    assert cluster.resource_yaml == expected_rc
 
 
 def test_get_local_queue_exists_fail(mocker):

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import string
+import sys
 from codeflare_sdk.ray.cluster.cluster import (
     Cluster,
     ClusterConfiguration,
@@ -255,6 +257,32 @@ def arg_check_del_effect(group, version, namespace, plural, name, *args):
         assert name == "ray-dashboard-unit-test-cluster-ray"
 
 
+def apply_template(yaml_file_path, variables):
+    with open(yaml_file_path, "r") as file:
+        yaml_content = file.read()
+
+    # Create a Template instance and substitute the variables
+    template = string.Template(yaml_content)
+    filled_yaml = template.substitute(variables)
+
+    # Now load the filled YAML into a Python object
+    return yaml.load(filled_yaml, Loader=yaml.FullLoader)
+
+
+def get_expected_image():
+    python_version = sys.version_info
+    if python_version.major == 3 and python_version.minor == 9:
+        return "quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06"
+    else:
+        return "quay.io/modh/ray@sha256:db667df1bc437a7b0965e8031e905d3ab04b86390d764d120e05ea5a5c18d1b4"
+
+
+def get_template_variables():
+    return {
+        "image": get_expected_image(),
+    }
+
+
 def arg_check_apply_effect(group, version, namespace, plural, body, *args):
     assert namespace == "ns"
     assert args == tuple()

src/codeflare_sdk/ray/cluster/test_config.py

@@ -13,15 +13,15 @@
 # limitations under the License.
 
 from codeflare_sdk.common.utils.unit_test_support import (
+    apply_template,
     createClusterWrongType,
-    get_local_queue,
     create_cluster_all_config_params,
+    get_template_variables,
 )
 from codeflare_sdk.ray.cluster.cluster import ClusterConfiguration, Cluster
 from pathlib import Path
 import filecmp
 import pytest
-import yaml
 import os
 
 parent = Path(__file__).resolve().parents[4]  # project directory
@@ -36,9 +36,11 @@ def test_default_cluster_creation(mocker):
 
     cluster = Cluster(ClusterConfiguration(name="default-cluster", namespace="ns"))
 
-    with open(f"{expected_clusters_dir}/ray/default-ray-cluster.yaml") as f:
-        expected_rc = yaml.load(f, Loader=yaml.FullLoader)
-        assert cluster.resource_yaml == expected_rc
+    expected_rc = apply_template(
+        f"{expected_clusters_dir}/ray/default-ray-cluster.yaml",
+        get_template_variables(),
+    )
+    assert cluster.resource_yaml == expected_rc
 
 
 def test_default_appwrapper_creation(mocker):
@@ -50,9 +52,10 @@ def test_default_appwrapper_creation(mocker):
         ClusterConfiguration(name="default-appwrapper", namespace="ns", appwrapper=True)
     )
 
-    with open(f"{expected_clusters_dir}/ray/default-appwrapper.yaml") as f:
-        expected_aw = yaml.load(f, Loader=yaml.FullLoader)
-        assert cluster.resource_yaml == expected_aw
+    expected_aw = apply_template(
+        f"{expected_clusters_dir}/ray/default-appwrapper.yaml", get_template_variables()
+    )
+    assert cluster.resource_yaml == expected_aw
 
 
 def test_config_creation_all_parameters(mocker):

tests/test_cluster_yamls/kueue/aw_kueue.yaml

@@ -38,7 +38,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:
@@ -103,7 +103,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:

tests/test_cluster_yamls/kueue/ray_cluster_kueue.yaml

@@ -38,7 +38,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:
@@ -103,7 +103,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:

tests/test_cluster_yamls/ray/default-appwrapper.yaml

@@ -36,7 +36,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:
@@ -101,7 +101,7 @@ spec:
           template:
             spec:
               containers:
-              - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+              - image: "${image}"
                 imagePullPolicy: Always
                 lifecycle:
                   preStop:

tests/test_cluster_yamls/ray/default-ray-cluster.yaml

@@ -28,7 +28,7 @@ spec:
     template:
       spec:
         containers:
-        - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+        - image: "${image}"
           imagePullPolicy: Always
           lifecycle:
             preStop:
@@ -93,7 +93,7 @@ spec:
     template:
       spec:
         containers:
-        - image: quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06
+        - image: "${image}"
           imagePullPolicy: Always
           lifecycle:
             preStop: