update tests

tests/test_onnx.py

@@ -1,8 +1,6 @@
 from pytorch_ner.onnx import onnx_export_and_check
-from pytorch_ner.utils import mkdir
 from tests.test_nn_modules.test_architecture import model_bilstm as model
 
 path_to_save = "models/model.onnx"
-mkdir("models")
 
 onnx_export_and_check(model=model, path_to_save=path_to_save)

tests/test_save.py

@@ -7,11 +7,9 @@
 from pytorch_ner.utils import rmdir
 from tests.test_train import label2idx, model, token2idx
 
-path_to_folder = "models/test_save/"
-path_to_onnx_folder = "models/test_onnx_save/"
-
-rmdir(path_to_folder)
-rmdir(path_to_onnx_folder)
+path_to_save_folder = "models/test"
+path_to_no_onnx_folder = os.path.join(path_to_save_folder, "no_onnx")
+path_to_onnx_folder = os.path.join(path_to_save_folder, "onnx")
 
 
 with open("config.yaml", "r") as fp:
@@ -20,7 +18,7 @@
 
 # without onnx
 save_model(
-    path_to_folder=path_to_folder,
+    path_to_folder=path_to_no_onnx_folder,
     model=model,
     token2idx=token2idx,
     label2idx=label2idx,
@@ -40,14 +38,14 @@
 
 
 class TestSave(unittest.TestCase):
-    def test_mkdir(self):
-        self.assertTrue(os.path.exists(path_to_folder))
-
     def test_num_files(self):
-        self.assertTrue(len(os.listdir(path_to_folder)) == 4)
+        self.assertTrue(len(os.listdir(os.listdir(path_to_no_onnx_folder)[0])) == 4)
 
     def test_num_files_with_onnx(self):
-        self.assertTrue(len(os.listdir(path_to_onnx_folder)) == 5)
+        self.assertTrue(len(os.listdir(os.listdir(path_to_onnx_folder)[0])) == 5)
+
+    def tearDownClass(cls):
+        rmdir(path_to_save_folder)
 
 
 if __name__ == "__main__":