6666 export_env ,
6767)
6868
69+ from tests .unit_test_support import (
70+ createClusterWithConfig ,
71+ createTestDDP ,
72+ createDDPJob_no_cluster ,
73+ createClusterConfig ,
74+ createDDPJob_with_cluster ,
75+ )
76+
6977import openshift
7078from openshift .selector import Selector
7179import ray
@@ -216,20 +224,7 @@ def test_auth_coverage():
216224
217225
218226def test_config_creation ():
219- config = ClusterConfiguration (
220- name = "unit-test-cluster" ,
221- namespace = "ns" ,
222- num_workers = 2 ,
223- min_cpus = 3 ,
224- max_cpus = 4 ,
225- min_memory = 5 ,
226- max_memory = 6 ,
227- num_gpus = 7 ,
228- instascale = True ,
229- machine_types = ["cpu.small" , "gpu.large" ],
230- image_pull_secrets = ["unit-test-pull-secret" ],
231- dispatch_priority = "default" ,
232- )
227+ config = createClusterConfig ()
233228
234229 assert config .name == "unit-test-cluster" and config .namespace == "ns"
235230 assert config .num_workers == 2
@@ -242,17 +237,15 @@ def test_config_creation():
242237 assert config .machine_types == ["cpu.small" , "gpu.large" ]
243238 assert config .image_pull_secrets == ["unit-test-pull-secret" ]
244239 assert config .dispatch_priority == "default"
245- return config
246240
247241
248242def test_cluster_creation ():
249- cluster = Cluster ( test_config_creation () )
243+ cluster = createClusterWithConfig ( )
250244 assert cluster .app_wrapper_yaml == "unit-test-cluster.yaml"
251245 assert cluster .app_wrapper_name == "unit-test-cluster"
252246 assert filecmp .cmp (
253247 "unit-test-cluster.yaml" , f"{ parent } /tests/test-case.yaml" , shallow = True
254248 )
255- return cluster
256249
257250
258251def test_default_cluster_creation (mocker ):
@@ -269,8 +262,6 @@ def test_default_cluster_creation(mocker):
269262 assert cluster .app_wrapper_name == "unit-test-default-cluster"
270263 assert cluster .config .namespace == "opendatahub"
271264
272- return cluster
273-
274265
275266def arg_check_apply_effect (group , version , namespace , plural , body , * args ):
276267 assert group == "mcad.ibm.com"
@@ -306,7 +297,7 @@ def test_cluster_up_down(mocker):
306297 "kubernetes.client.CustomObjectsApi.list_cluster_custom_object" ,
307298 return_value = {"items" : []},
308299 )
309- cluster = test_cluster_creation ()
300+ cluster = cluster = createClusterWithConfig ()
310301 cluster .up ()
311302 cluster .down ()
312303
@@ -374,7 +365,7 @@ def test_cluster_uris(mocker):
374365 side_effect = uri_retreival ,
375366 )
376367
377- cluster = test_cluster_creation ()
368+ cluster = cluster = createClusterWithConfig ()
378369 assert cluster .cluster_uri () == "ray://unit-test-cluster-head-svc.ns.svc:10001"
379370 assert (
380371 cluster .cluster_dashboard_uri ()
@@ -421,7 +412,7 @@ def test_ray_job_wrapping(mocker):
421412 "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object" ,
422413 side_effect = uri_retreival ,
423414 )
424- cluster = test_cluster_creation ()
415+ cluster = cluster = createClusterWithConfig ()
425416
426417 mocker .patch (
427418 "ray.job_submission.JobSubmissionClient._check_connection_and_version_with_url" ,
@@ -1771,7 +1762,7 @@ def test_wait_ready(mocker, capsys):
17711762
17721763def test_jobdefinition_coverage ():
17731764 abstract = JobDefinition ()
1774- cluster = Cluster ( test_config_creation () )
1765+ cluster = createClusterWithConfig ( )
17751766 abstract ._dry_run (cluster )
17761767 abstract .submit (cluster )
17771768
@@ -1783,22 +1774,7 @@ def test_job_coverage():
17831774
17841775
17851776def test_DDPJobDefinition_creation ():
1786- ddp = DDPJobDefinition (
1787- script = "test.py" ,
1788- m = None ,
1789- script_args = ["test" ],
1790- name = "test" ,
1791- cpu = 1 ,
1792- gpu = 0 ,
1793- memMB = 1024 ,
1794- h = None ,
1795- j = "2x1" ,
1796- env = {"test" : "test" },
1797- max_retries = 0 ,
1798- mounts = [],
1799- rdzv_port = 29500 ,
1800- scheduler_args = {"requirements" : "test" },
1801- )
1777+ ddp = createTestDDP ()
18021778 assert ddp .script == "test.py"
18031779 assert ddp .m == None
18041780 assert ddp .script_args == ["test" ]
@@ -1813,7 +1789,6 @@ def test_DDPJobDefinition_creation():
18131789 assert ddp .mounts == []
18141790 assert ddp .rdzv_port == 29500
18151791 assert ddp .scheduler_args == {"requirements" : "test" }
1816- return ddp
18171792
18181793
18191794def test_DDPJobDefinition_dry_run (mocker ):
@@ -1827,8 +1802,8 @@ def test_DDPJobDefinition_dry_run(mocker):
18271802 "codeflare_sdk.cluster.cluster.Cluster.cluster_dashboard_uri" ,
18281803 return_value = "" ,
18291804 )
1830- ddp = test_DDPJobDefinition_creation ()
1831- cluster = Cluster ( test_config_creation () )
1805+ ddp = createTestDDP ()
1806+ cluster = createClusterWithConfig ( )
18321807 ddp_job = ddp ._dry_run (cluster )
18331808 assert type (ddp_job ) == AppDryRunInfo
18341809 assert ddp_job ._fmt is not None
@@ -1863,7 +1838,7 @@ def test_DDPJobDefinition_dry_run_no_cluster(mocker):
18631838 return_value = "opendatahub" ,
18641839 )
18651840
1866- ddp = test_DDPJobDefinition_creation ()
1841+ ddp = createTestDDP ()
18671842 ddp .image = "fake-image"
18681843 ddp_job = ddp ._dry_run_no_cluster ()
18691844 assert type (ddp_job ) == AppDryRunInfo
@@ -1900,7 +1875,7 @@ def test_DDPJobDefinition_dry_run_no_resource_args(mocker):
19001875 "codeflare_sdk.cluster.cluster.Cluster.cluster_dashboard_uri" ,
19011876 return_value = "" ,
19021877 )
1903- cluster = Cluster ( test_config_creation () )
1878+ cluster = createClusterWithConfig ( )
19041879 ddp = DDPJobDefinition (
19051880 script = "test.py" ,
19061881 m = None ,
@@ -1936,7 +1911,7 @@ def test_DDPJobDefinition_dry_run_no_cluster_no_resource_args(mocker):
19361911 return_value = "opendatahub" ,
19371912 )
19381913
1939- ddp = test_DDPJobDefinition_creation ()
1914+ ddp = createTestDDP ()
19401915 try :
19411916 ddp ._dry_run_no_cluster ()
19421917 assert 0 == 1
@@ -1988,8 +1963,8 @@ def test_DDPJobDefinition_submit(mocker):
19881963 "codeflare_sdk.cluster.cluster.Cluster.cluster_dashboard_uri" ,
19891964 return_value = "fake-dashboard-uri" ,
19901965 )
1991- ddp_def = test_DDPJobDefinition_creation ()
1992- cluster = Cluster ( test_config_creation () )
1966+ ddp_def = createTestDDP ()
1967+ cluster = createClusterWithConfig ( )
19931968 mocker .patch (
19941969 "codeflare_sdk.job.jobs.get_current_namespace" ,
19951970 side_effect = "opendatahub" ,
@@ -2019,13 +1994,13 @@ def test_DDPJob_creation(mocker):
20191994 "codeflare_sdk.cluster.cluster.Cluster.cluster_dashboard_uri" ,
20201995 return_value = "fake-dashboard-uri" ,
20211996 )
2022- ddp_def = test_DDPJobDefinition_creation ()
2023- cluster = Cluster ( test_config_creation () )
1997+ ddp_def = createTestDDP ()
1998+ cluster = createClusterWithConfig ( )
20241999 mocker .patch (
20252000 "codeflare_sdk.job.jobs.torchx_runner.schedule" ,
20262001 return_value = "fake-dashboard-url" ,
20272002 ) # a fake app_handle
2028- ddp_job = DDPJob (ddp_def , cluster )
2003+ ddp_job = createDDPJob_with_cluster (ddp_def , cluster )
20292004 assert type (ddp_job ) == DDPJob
20302005 assert type (ddp_job .job_definition ) == DDPJobDefinition
20312006 assert type (ddp_job .cluster ) == Cluster
@@ -2038,11 +2013,10 @@ def test_DDPJob_creation(mocker):
20382013 assert type (job_info ._app ) == AppDef
20392014 assert type (job_info ._cfg ) == type (dict ())
20402015 assert type (job_info ._scheduler ) == type (str ())
2041- return ddp_job
20422016
20432017
20442018def test_DDPJob_creation_no_cluster (mocker ):
2045- ddp_def = test_DDPJobDefinition_creation ()
2019+ ddp_def = createTestDDP ()
20462020 ddp_def .image = "fake-image"
20472021 mocker .patch (
20482022 "codeflare_sdk.job.jobs.get_current_namespace" ,
@@ -2052,7 +2026,7 @@ def test_DDPJob_creation_no_cluster(mocker):
20522026 "codeflare_sdk.job.jobs.torchx_runner.schedule" ,
20532027 return_value = "fake-app-handle" ,
20542028 ) # a fake app_handle
2055- ddp_job = DDPJob (ddp_def , None )
2029+ ddp_job = createDDPJob_no_cluster (ddp_def , None )
20562030 assert type (ddp_job ) == DDPJob
20572031 assert type (ddp_job .job_definition ) == DDPJobDefinition
20582032 assert ddp_job .cluster == None
@@ -2065,11 +2039,14 @@ def test_DDPJob_creation_no_cluster(mocker):
20652039 assert type (job_info ._app ) == AppDef
20662040 assert type (job_info ._cfg ) == type (dict ())
20672041 assert type (job_info ._scheduler ) == type (str ())
2068- return ddp_job
20692042
20702043
20712044def test_DDPJob_status (mocker ):
2072- ddp_job = test_DDPJob_creation (mocker )
2045+ # Setup the neccesary mock patches
2046+ test_DDPJob_creation (mocker )
2047+ ddp_def = createTestDDP ()
2048+ cluster = createClusterWithConfig ()
2049+ ddp_job = createDDPJob_with_cluster (ddp_def , cluster )
20732050 mocker .patch (
20742051 "codeflare_sdk.job.jobs.torchx_runner.status" , return_value = "fake-status"
20752052 )
@@ -2079,7 +2056,11 @@ def test_DDPJob_status(mocker):
20792056
20802057
20812058def test_DDPJob_logs (mocker ):
2082- ddp_job = test_DDPJob_creation (mocker )
2059+ # Setup the neccesary mock patches
2060+ test_DDPJob_creation (mocker )
2061+ ddp_def = createTestDDP ()
2062+ cluster = createClusterWithConfig ()
2063+ ddp_job = createDDPJob_with_cluster (ddp_def , cluster )
20832064 mocker .patch (
20842065 "codeflare_sdk.job.jobs.torchx_runner.log_lines" , return_value = "fake-logs"
20852066 )
@@ -2093,7 +2074,11 @@ def arg_check_side_effect(*args):
20932074
20942075
20952076def test_DDPJob_cancel (mocker ):
2096- ddp_job = test_DDPJob_creation_no_cluster (mocker )
2077+ # Setup the neccesary mock patches
2078+ test_DDPJob_creation_no_cluster (mocker )
2079+ ddp_def = createTestDDP ()
2080+ ddp_def .image = "fake-image"
2081+ ddp_job = createDDPJob_no_cluster (ddp_def , None )
20972082 mocker .patch (
20982083 "openshift.get_project_name" ,
20992084 return_value = "opendatahub" ,
0 commit comments