Fix dataset tests

mlcommons · Sep 13, 2023 · 146bb19 · 146bb19
1 parent afaedb8
commit 146bb19
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 18 deletions.
diff --git a/cli/medperf/commands/dataset/create.py b/cli/medperf/commands/dataset/create.py
@@ -190,7 +190,6 @@ def run_sanity_check(self):
         }
         sanity_str_params = {
             "Ptasks.sanity_check.parameters.input.data_path.opts": "ro",
-            "Ptasks.sanity_check.parameters.input.report_file.opts": "ro",
         }
 
         if self.labels_specified:
@@ -199,6 +198,9 @@ def run_sanity_check(self):
 
         if self.report_specified:
             sanity_params["report_file"] = out_report
+            sanity_str_params[
+                "Ptasks.sanity_check.parameters.input.report_file.opts"
+            ] = "ro"
 
         self.ui.text = "Running sanity check..."
         try:
@@ -220,9 +222,11 @@ def run_sanity_check(self):
     def run_statistics(self):
         statistics_timeout = config.statistics_timeout
         out_datapath = self.out_datapath
+        out_labelspath = self.out_labelspath
 
         statistics_params = {
             "data_path": out_datapath,
+            "labels_path": out_labelspath,
             "output_path": self.out_statistics_path,
         }
         statistics_str_params = {

diff --git a/cli/medperf/tests/commands/dataset/test_create.py b/cli/medperf/tests/commands/dataset/test_create.py
@@ -21,14 +21,15 @@
 NAME = "name"
 DESCRIPTION = "description"
 LOCATION = "location"
+SUMMARY_PATH = "summary"
+REPORT_PATH = "report"
 
 
 @pytest.fixture
 def preparation(mocker, comms, ui):
     mocker.patch("os.path.abspath", side_effect=lambda x: x)
     mocker.patch(
-        PATCH_DATAPREP.format("generate_tmp_path"),
-        side_effect=[OUT_PATH, STATISTICS_PATH],
+        PATCH_DATAPREP.format("generate_tmp_path"), return_value=STATISTICS_PATH
     )
     mocker.patch(PATCH_DATAPREP.format("Benchmark.get"), return_value=TestBenchmark())
     preparation = DataPreparation(
@@ -39,11 +40,17 @@ def preparation(mocker, comms, ui):
         NAME,
         DESCRIPTION,
         LOCATION,
+        SUMMARY_PATH,
     )
     mocker.patch(PATCH_DATAPREP.format("Cube.get"), return_value=MockCube(True))
     preparation.get_prep_cube()
     preparation.data_path = DATA_PATH
     preparation.labels_path = LABELS_PATH
+    preparation.out_datapath = OUT_DATAPATH
+    preparation.out_labelspath = OUT_LABELSPATH
+    preparation.report_path = REPORT_PATH
+    preparation.report_specified = False
+    preparation.labels_specified = True
     return preparation
 
 
@@ -81,7 +88,7 @@ def test_get_prep_cube_gets_prep_cube_if_provided(
         )
 
         # Act
-        preparation = DataPreparation(None, cube_uid, *[""] * 5)
+        preparation = DataPreparation(None, cube_uid, *[""] * 6)
         preparation.get_prep_cube()
 
         # Assert
@@ -99,7 +106,7 @@ def test_get_prep_cube_gets_benchmark_cube_if_provided(
         )
 
         # Act
-        preparation = DataPreparation(cube_uid, None, *[""] * 5)
+        preparation = DataPreparation(cube_uid, None, *[""] * 6)
         preparation.get_prep_cube()
 
         # Assert
@@ -139,19 +146,29 @@ def test_run_cube_tasks_runs_required_tasks(self, mocker, preparation):
         calls = [prepare, check, stats]
 
         # Act
-        preparation.run_cube_tasks()
+        preparation.run_prepare()
+        preparation.run_sanity_check()
+        preparation.run_statistics()
 
         # Assert
         spy.assert_has_calls(calls)
 
     def test_run_executes_expected_flow(self, mocker, comms, ui):
         # Arrange
         validate_spy = mocker.patch(PATCH_DATAPREP.format("DataPreparation.validate"))
-        get_cube_spy = mocker.patch(
-            PATCH_DATAPREP.format("DataPreparation.get_prep_cube")
+        get_cube_spy = mocker.spy(DataPreparation, "get_prep_cube")
+        mocker.patch(
+            PATCH_DATAPREP.format("Cube.get"),
+            side_effect=lambda id: MockCube(True, id),
+        )
+        run_prepare_spy = mocker.patch(
+            PATCH_DATAPREP.format("DataPreparation.run_prepare")
+        )
+        run_sanity_check_spy = mocker.patch(
+            PATCH_DATAPREP.format("DataPreparation.run_sanity_check")
         )
-        run_cube_spy = mocker.patch(
-            PATCH_DATAPREP.format("DataPreparation.run_cube_tasks")
+        run_statistics_spy = mocker.patch(
+            PATCH_DATAPREP.format("DataPreparation.run_statistics")
         )
         get_stat_spy = mocker.patch(
             PATCH_DATAPREP.format("DataPreparation.get_statistics"),
@@ -172,7 +189,9 @@ def test_run_executes_expected_flow(self, mocker, comms, ui):
         # Assert
         validate_spy.assert_called_once()
         get_cube_spy.assert_called_once()
-        run_cube_spy.assert_called_once()
+        run_prepare_spy.assert_called_once()
+        run_sanity_check_spy.assert_called_once()
+        run_statistics_spy.assert_called_once()
         get_stat_spy.assert_called_once()
         generate_uids_spy.assert_called_once()
         to_permanent_path_spy.assert_called_once()
@@ -185,7 +204,7 @@ def test_fails_if_invalid_params(self, mocker, benchmark_uid, cube_uid, comms, u
         num_arguments = int(benchmark_uid is None) + int(cube_uid is None)
 
         # Act
-        preparation = DataPreparation(benchmark_uid, cube_uid, *[""] * 5)
+        preparation = DataPreparation(benchmark_uid, cube_uid, *[""] * 6)
         # Assert
 
         if num_arguments != 1:
@@ -288,8 +307,10 @@ def generate_uids(cls):
         cls.generated_uid = uid
 
     mocker.patch(PATCH_DATAPREP.format("DataPreparation.validate"))
-    mocker.patch(PATCH_DATAPREP.format("DataPreparation.get_prep_cube"))
-    mocker.patch(PATCH_DATAPREP.format("DataPreparation.run_cube_tasks"))
+    mocker.patch(PATCH_DATAPREP.format("DataPreparation.run_prepare"))
+    mocker.patch(PATCH_DATAPREP.format("DataPreparation.run_sanity_check"))
+    mocker.patch(PATCH_DATAPREP.format("DataPreparation.run_statistics"))
+    mocker.patch(PATCH_DATAPREP.format("SummaryGenerator.run"))
     mocker.patch(
         PATCH_DATAPREP.format("DataPreparation.get_statistics"),
     )
@@ -304,9 +325,13 @@ def generate_uids(cls):
     mocker.patch(
         PATCH_DATAPREP.format("DataPreparation.write"),
     )
+    mocker.patch(
+        PATCH_DATAPREP.format("Cube.get"),
+        side_effect=lambda id: MockCube(True, id),
+    )
 
     # Act
-    returned_uid = DataPreparation.run("", "", "", "")
+    returned_uid = DataPreparation.run("", 1, "", "")
 
     # Assert
     assert returned_uid == uid
diff --git a/cli/medperf/tests/commands/result/test_create.py b/cli/medperf/tests/commands/result/test_create.py
@@ -323,6 +323,6 @@ def test_execution_of_one_model_writes_result(self, mocker, setup):
 
         # Assert
         assert (
-            yaml.load(open(expected_file))["results"]
+            yaml.safe_load(open(expected_file))["results"]
             == self.state_variables["models_props"][model_uid]["results"]
         )
diff --git a/cli/medperf/tests/mocks/cube.py b/cli/medperf/tests/mocks/cube.py
@@ -6,10 +6,11 @@
 
 
 class MockCube:
-    def __init__(self, is_valid):
+    def __init__(self, is_valid, id=1, report_specified=False):
         self.name = "Test"
         self.is_valid = is_valid
-        self.id = 1
+        self.id = id
+        self.report_specified = report_specified
 
     def valid(self):
         return self.is_valid
@@ -18,6 +19,8 @@ def run(self):
         pass
 
     def get_default_output(self, *args, **kwargs):
+        if args == ("prepare", "report_file") and not self.report_specified:
+            return None
         return "out_path"
 
     @property