8185 - Refactor test

tests/test_bundle_ckpt_export.py

@@ -19,9 +19,10 @@
 from parameterized import parameterized
 
 from monai.bundle import ConfigParser
+from monai.bundle.scripts import ckpt_export
 from monai.data import load_net_with_metadata
 from monai.networks import save_state
-from tests.utils import command_line_tests, skip_if_windows
+from tests.utils import skip_if_windows
 
 TEST_CASE_1 = ["", ""]
 
@@ -37,67 +38,69 @@ def setUp(self):
         self.device = os.environ.get("CUDA_VISIBLE_DEVICES")
         if not self.device:
             os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # default
+        self.meta_file = os.path.join(os.path.dirname(__file__), "testing_data", "metadata.json")
+        self.config_file = os.path.join(os.path.dirname(__file__), "testing_data", "inference.json")
+        self.tempdir_obj = tempfile.TemporaryDirectory()
+        tempdir = self.tempdir_obj.name
+        self.def_args = {"meta_file": "will be replaced by `meta_file` arg"}
+        self.def_args_file = os.path.join(tempdir, "def_args.yaml")
+
+        self.ckpt_file = os.path.join(tempdir, "model.pt")
+        self.ts_file = os.path.join(tempdir, "model.ts")
+
+        self.parser = ConfigParser()
+        self.parser.export_config_file(config=self.def_args, filepath=self.def_args_file)
+        self.parser.read_config(self.config_file)
+        self.net = self.parser.get_parsed_content("network_def")
 
     def tearDown(self):
         if self.device is not None:
             os.environ["CUDA_VISIBLE_DEVICES"] = self.device
         else:
             del os.environ["CUDA_VISIBLE_DEVICES"]  # previously unset
+        self.tempdir_obj.cleanup()
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_export(self, key_in_ckpt, use_trace):
-        meta_file = os.path.join(os.path.dirname(__file__), "testing_data", "metadata.json")
-        config_file = os.path.join(os.path.dirname(__file__), "testing_data", "inference.json")
-        with tempfile.TemporaryDirectory() as tempdir:
-            def_args = {"meta_file": "will be replaced by `meta_file` arg"}
-            def_args_file = os.path.join(tempdir, "def_args.yaml")
-
-            ckpt_file = os.path.join(tempdir, "model.pt")
-            ts_file = os.path.join(tempdir, "model.ts")
-
-            parser = ConfigParser()
-            parser.export_config_file(config=def_args, filepath=def_args_file)
-            parser.read_config(config_file)
-            net = parser.get_parsed_content("network_def")
-            save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file)
-
-            cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "network_def", "--filepath", ts_file]
-            cmd += ["--meta_file", meta_file, "--config_file", f"['{config_file}','{def_args_file}']", "--ckpt_file"]
-            cmd += [ckpt_file, "--key_in_ckpt", key_in_ckpt, "--args_file", def_args_file]
-            if use_trace == "True":
-                cmd += ["--use_trace", use_trace, "--input_shape", "[1, 1, 96, 96, 96]"]
-            command_line_tests(cmd)
-            self.assertTrue(os.path.exists(ts_file))
-
-            _, metadata, extra_files = load_net_with_metadata(
-                ts_file, more_extra_files=["inference.json", "def_args.json"]
-            )
-            self.assertIn("schema", metadata)
-            self.assertIn("meta_file", json.loads(extra_files["def_args.json"]))
-            self.assertIn("network_def", json.loads(extra_files["inference.json"]))
+    def test_ckpt_export_default(self, key_in_ckpt, use_trace):
+        ckpt_file = os.path.join(self.tempdir_obj.name, "models/model.pt")
+        ts_file = os.path.join(self.tempdir_obj.name, "models/model.ts")
+
+        save_state(src=self.net if key_in_ckpt == "" else {key_in_ckpt: self.net}, path=ckpt_file)
+        ckpt_export(
+            net_id="network_def",
+            filepath=ts_file,
+            meta_file=self.meta_file,
+            config_file=self.config_file,
+            ckpt_file=ckpt_file,
+            key_in_ckpt=key_in_ckpt,
+            args_file=self.def_args_file,
+            use_trace=use_trace,
+            input_shape=[1, 1, 96, 96, 96] if use_trace == "True" else None,
+        )
+        self.assertTrue(os.path.exists(ts_file))
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_default_value(self, key_in_ckpt, use_trace):
-        config_file = os.path.join(os.path.dirname(__file__), "testing_data", "inference.json")
-        with tempfile.TemporaryDirectory() as tempdir:
-            def_args = {"meta_file": "will be replaced by `meta_file` arg"}
-            def_args_file = os.path.join(tempdir, "def_args.yaml")
-            ckpt_file = os.path.join(tempdir, "models/model.pt")
-            ts_file = os.path.join(tempdir, "models/model.ts")
-
-            parser = ConfigParser()
-            parser.export_config_file(config=def_args, filepath=def_args_file)
-            parser.read_config(config_file)
-            net = parser.get_parsed_content("network_def")
-            save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file)
-
-            # check with default value
-            cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "--key_in_ckpt", key_in_ckpt]
-            cmd += ["--config_file", config_file, "--bundle_root", tempdir]
-            if use_trace == "True":
-                cmd += ["--use_trace", use_trace, "--input_shape", "[1, 1, 96, 96, 96]"]
-            command_line_tests(cmd)
-            self.assertTrue(os.path.exists(ts_file))
+    def test_ckpt_export(self, key_in_ckpt, use_trace):
+        save_state(src=self.net if key_in_ckpt == "" else {key_in_ckpt: self.net}, path=self.ckpt_file)
+        ckpt_export(
+            net_id="network_def",
+            filepath=self.ts_file,
+            meta_file=self.meta_file,
+            config_file=[self.config_file, self.def_args_file],
+            ckpt_file=self.ckpt_file,
+            key_in_ckpt=key_in_ckpt,
+            args_file=self.def_args_file,
+            use_trace=use_trace,
+            input_shape=[1, 1, 96, 96, 96] if use_trace == "True" else None,
+        )
+        self.assertTrue(os.path.exists(self.ts_file))
+
+        _, metadata, extra_files = load_net_with_metadata(
+            self.ts_file, more_extra_files=["inference.json", "def_args.json"]
+        )
+        self.assertIn("schema", metadata)
+        self.assertIn("meta_file", json.loads(extra_files["def_args.json"]))
+        self.assertIn("network_def", json.loads(extra_files["inference.json"]))
 
 
 if __name__ == "__main__":

tests/test_clip_intensity_percentilesd.py

@@ -18,10 +18,9 @@
 from monai.transforms import ClipIntensityPercentilesd
 from monai.transforms.utils_pytorch_numpy_unification import clip, percentile
 from monai.utils.type_conversion import convert_to_tensor
+from tests.test_clip_intensity_percentiles import test_hard_clip_func, test_soft_clip_func
 from tests.utils import TEST_NDARRAYS, NumpyImageTestCase2D, NumpyImageTestCase3D, assert_allclose
 
-from .test_clip_intensity_percentiles import test_hard_clip_func, test_soft_clip_func
-
 
 class TestClipIntensityPercentilesd2D(NumpyImageTestCase2D):