diff --git a/aidge_onnx/unit_tests/test_import_export.py b/aidge_onnx/unit_tests/test_import_export.py index d40358a49ac13dccb391dcf8be468fab1d62c0e0..e850a1654b6fd3f0d4efb216c24d867ae40096a8 100644 --- a/aidge_onnx/unit_tests/test_import_export.py +++ b/aidge_onnx/unit_tests/test_import_export.py @@ -26,7 +26,7 @@ class ImportExportONNXTestCases: def download_onnx_model(url, destination): - if os.path.exists(destination) : + if os.path.exists(destination): return try: response = requests.get(url, stream=True) @@ -51,17 +51,14 @@ class TestLoadSave(unittest.TestCase): input_dims=np.array([[1, 3, 224, 224]]), output_dims=np.array([[1, 1000]]), ), - # ImportExportONNXTestCases( - # url="https://github.com/onnx/models/raw/main/validated/vision/classification/squeezenet/model/squeezenet1.1-7.onnx?download=", - # original_onnx_filename="squeezenet1.1-7.onnx", - # input_dims=np.array([[1, 3, 224, 224]]), - # output_dims=np.array([[1, 1000, 1, 1]]), - # ), ] return def tearDown(self): - pass + for test in self.test_runs: + aidge_onnx_file = "aidge_" + test.original_onnx_filename + if os.path.exists(aidge_onnx_file): + os.remove(aidge_onnx_file) def test_import_export(self): Log.set_console_level(Level.Info) @@ -71,7 +68,7 @@ class TestLoadSave(unittest.TestCase): download_onnx_model(test.url, test.original_onnx_filename) Log.notice("\nConverting onnx model to aidge") model = aidge_onnx.load_onnx(test.original_onnx_filename, verbose=True) - aidge_model_onnx_filename = "aidge" + test.original_onnx_filename + aidge_model_onnx_filename = "aidge_" + test.original_onnx_filename Log.notice("\nRe exporting aidge model to onnx...") aidge_onnx.export_onnx( @@ -86,6 +83,5 @@ class TestLoadSave(unittest.TestCase): self.assertTrue(aidge_onnx.check_onnx_validity(aidge_model_onnx_filename)) - if __name__ == "__main__": unittest.main()