Directly tests export_ckpt function instead of using command_line_tests

Performance Before:
94.81s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_default_value_1_model
20.95s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_default_value_0_
15.26s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_export_2_model
14.86s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_default_value_2_model
14.55s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_export_1_model
14.28s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_export_0_

Performance after:

1.62s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_2_model
1.25s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_default_2_model
0.64s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_0_
0.57s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_1_model
0.57s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_default_1_model
0.55s call     tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_default_0_
0.01s setup    tests/test_bundle_ckpt_export.py::TestCKPTExport::test_ckpt_export_0_

Author: garciadias

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 = ["", ""]
 
@@ -51,8 +52,6 @@ def setUp(self):
         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")
-        self.cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "network_def", "--filepath", self.ts_file]
-        self.cmd += ["--meta_file", self.meta_file, "--config_file", f"['{self.config_file}','{self.def_args_file}']", "--ckpt_file"]
 
     def tearDown(self):
         if self.device is not None:
@@ -62,36 +61,47 @@ def tearDown(self):
         self.tempdir_obj.cleanup()
 
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_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)
-            full_cmd = self.cmd + [self.ckpt_file, "--key_in_ckpt", key_in_ckpt, "--args_file", self.def_args_file]
-            if use_trace == "True":
-                full_cmd += ["--use_trace", use_trace, "--input_shape", "[1, 1, 96, 96, 96]"]
-            command_line_tests(full_cmd)
-            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"]))
-
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
-    def test_default_value(self, key_in_ckpt, use_trace):
+    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)
-
-        # check with default value
-        cmd = ["coverage", "run", "-m", "monai.bundle", "ckpt_export", "--key_in_ckpt", key_in_ckpt]
-        cmd += ["--config_file", self.config_file, "--bundle_root", self.tempdir_obj.name]
-        if use_trace == "True":
-            cmd += ["--use_trace", use_trace, "--input_shape", "[1, 1, 96, 96, 96]"]
-        command_line_tests(cmd)
+        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_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__":
     unittest.main()