Include stages file to mlcube download tests

mlcommons · Oct 25, 2024 · 39cd511 · 39cd511
1 parent 4a146d3
commit 39cd511
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
diff --git a/cli/medperf/tests/entities/test_cube.py b/cli/medperf/tests/entities/test_cube.py
@@ -68,7 +68,8 @@ def set_common_attributes(self, setup):
             self.cube_path, config.additional_path, config.tarball_filename
         )
         self.img_path = os.path.join(self.cube_path, config.image_path, "img.tar.gz")
-        self.config_files_paths = [self.manifest_path, self.params_path]
+        self.stages_path = os.path.join(self.cube_path, config.stages_filename)
+        self.config_files_paths = [self.manifest_path, self.params_path, self.stages_path]
         self.run_files_paths = [self.add_path, self.img_path]
 
     @pytest.mark.parametrize("setup", [{"remote": [DEFAULT_CUBE]}], indirect=True)

diff --git a/cli/medperf/tests/entities/utils.py b/cli/medperf/tests/entities/utils.py
@@ -112,16 +112,21 @@ def setup_cube_comms_downloads(mocker, fs):
     add_file = config.tarball_filename
     img_path = config.image_path
     img_file = "img.tar.gz"
+    stages_path = ""
+    stages_file = config.stages_filename
 
     get_cube_fn = generate_cubefile_fn(fs, cube_path, cube_file)
     get_params_fn = generate_cubefile_fn(fs, params_path, params_file)
     get_add_fn = generate_cubefile_fn(fs, add_path, add_file)
     get_img_fn = generate_cubefile_fn(fs, img_path, img_file)
+    get_stages_fn = generate_cubefile_fn(fs, stages_path, stages_file)
+
 
     mocker.patch(PATCH_RESOURCES.format("get_cube"), side_effect=get_cube_fn)
     mocker.patch(PATCH_RESOURCES.format("get_cube_params"), side_effect=get_params_fn)
     mocker.patch(PATCH_RESOURCES.format("get_cube_additional"), side_effect=get_add_fn)
     mocker.patch(PATCH_RESOURCES.format("get_cube_image"), side_effect=get_img_fn)
+    mocker.patch(PATCH_RESOURCES.format("get_cube_stages"), side_effect=get_stages_fn)
 
 
 # Setup Dataset

diff --git a/cli/medperf/tests/mocks/cube.py b/cli/medperf/tests/mocks/cube.py
@@ -18,5 +18,7 @@ class TestCube(Cube):
         "https://test.com/additional_files.tar.gz"
     )
     additional_files_tarball_hash: Optional[str] = EMPTY_FILE_HASH
+    stages_url: Optional[str] = "https://test.com/stages.yaml"
+    stages_hash: Optional[str] = EMPTY_FILE_HASH
     state: str = "OPERATION"
     is_valid = True