Add pytest fixture for cleaning work directory

Author: ndem0

tests/conftest.py

@@ -0,0 +1,16 @@
+import shutil
+from pathlib import Path
+import pytest
+
+@pytest.fixture
+def clean_tmp_dir(tmp_path):
+    path = Path(tmp_path)
+
+    if path.exists():
+        shutil.rmtree(path)
+
+    path.mkdir(parents=True, exist_ok=True)
+    yield path
+
+    if path.exists():
+        shutil.rmtree(path)

tests/test_solver/test_competitive_pinn.py

@@ -113,9 +113,8 @@ def test_solver_test(problem, batch_size, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_train_load_restore(problem):
-    dir = "tests/test_solver/tmp"
-    problem = problem
+def test_train_load_restore(clean_tmp_dir, problem):
+    dir = clean_tmp_dir
     solver = CompPINN(problem=problem, model=model)
     trainer = Trainer(
         solver=solver,
@@ -150,9 +149,4 @@ def test_train_load_restore(problem):
     )
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")
+    )

tests/test_solver/test_ensemble_pinn.py

@@ -107,8 +107,8 @@ def test_solver_test(batch_size, compile):
         )
 
 
-def test_train_load_restore():
-    dir = "tests/test_solver/tmp"
+def test_train_load_restore(clean_tmp_dir):
+    dir = clean_tmp_dir
     solver = DeepEnsemblePINN(models=models, problem=problem)
     trainer = Trainer(
         solver=solver,
@@ -141,8 +141,3 @@ def test_train_load_restore():
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
     )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")

tests/test_solver/test_ensemble_supervised_solver.py

@@ -235,8 +235,8 @@ def test_solver_test_graph(batch_size, use_lt):
     trainer.test()
 
 
-def test_train_load_restore():
-    dir = "tests/test_solver/tmp/"
+def test_train_load_restore(clean_tmp_dir):
+    dir = clean_tmp_dir
     problem = LabelTensorProblem()
     solver = DeepEnsembleSupervisedSolver(problem=problem, models=models)
     trainer = Trainer(
@@ -270,8 +270,3 @@ def test_train_load_restore():
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
     )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")

tests/test_solver/test_garom.py

@@ -163,9 +163,8 @@ def test_solver_test(batch_size, compile):
         )
 
 
-def test_train_load_restore():
-    dir = "tests/test_solver/tmp/"
-    problem = TensorProblem()
+def test_train_load_restore(clean_tmp_dir):
+    dir = clean_tmp_dir
     solver = GAROM(
         problem=TensorProblem(),
         generator=Generator(),
@@ -201,8 +200,3 @@ def test_train_load_restore():
     test_pts = torch.rand(20, 1)
     assert new_solver.forward(test_pts).shape == (20, 2)
     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")

tests/test_solver/test_gradient_pinn.py

@@ -119,9 +119,8 @@ def test_solver_test(problem, batch_size, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_train_load_restore(problem):
-    dir = "tests/test_solver/tmp"
-    problem = problem
+def test_train_load_restore(clean_tmp_dir, problem):
+    dir = clean_tmp_dir
     solver = GradientPINN(model=model, problem=problem)
     trainer = Trainer(
         solver=solver,
@@ -157,8 +156,3 @@ def test_train_load_restore(problem):
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
     )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")

tests/test_solver/test_pinn.py

@@ -101,9 +101,8 @@ def test_solver_test(problem, batch_size, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_train_load_restore(problem):
-    dir = "tests/test_solver/tmp"
-    problem = problem
+def test_train_load_restore(clean_tmp_dir, problem):
+    dir = clean_tmp_dir
     solver = PINN(model=model, problem=problem)
     trainer = Trainer(
         solver=solver,
@@ -116,6 +115,8 @@ def test_train_load_restore(problem):
         default_root_dir=dir,
     )
     trainer.train()
+    import os
+    print(os.listdir(f"{dir}/lightning_logs/version_0/checkpoints/"))
 
     # restore
     new_trainer = Trainer(solver=solver, max_epochs=5, accelerator="cpu")
@@ -124,6 +125,7 @@ def test_train_load_restore(problem):
         + "epoch=4-step=5.ckpt"
     )
 
+
     # loading
     new_solver = PINN.load_from_checkpoint(
         f"{dir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt",
@@ -136,9 +138,4 @@ def test_train_load_restore(problem):
     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")
+    )

tests/test_solver/test_rba_pinn.py

@@ -122,9 +122,8 @@ def test_solver_test(problem, batch_size, loss, compile):
 
 
 @pytest.mark.parametrize("problem", [problem, inverse_problem])
-def test_train_load_restore(problem):
-    dir = "tests/test_solver/tmp"
-    problem = problem
+def test_train_load_restore(clean_tmp_dir, problem):
+    dir = clean_tmp_dir
     solver = RBAPINN(model=model, problem=problem)
     trainer = Trainer(
         solver=solver,
@@ -159,9 +158,4 @@ def test_train_load_restore(problem):
     )
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")
+    )

tests/test_solver/test_supervised_solver.py

@@ -212,8 +212,25 @@ def test_solver_test_graph(batch_size, use_lt):
     trainer.test()
 
 
-def test_train_load_restore():
-    dir = "tests/test_solver/tmp/"
+import shutil
+from pathlib import Path
+import pytest
+
+@pytest.fixture
+def clean_tmp_dir():
+    path = Path("tests/test_solver/tmp/")
+
+    if path.exists():
+        shutil.rmtree(path)
+
+    path.mkdir(parents=True, exist_ok=True)
+    yield path
+
+    if path.exists():
+        shutil.rmtree(path)
+
+def test_train_load_restore(clean_tmp_dir):
+    dir = clean_tmp_dir
     problem = LabelTensorProblem()
     solver = SupervisedSolver(problem=problem, model=model)
     trainer = Trainer(
@@ -247,9 +264,4 @@ def test_train_load_restore():
     assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape
     torch.testing.assert_close(
         new_solver.forward(test_pts), solver.forward(test_pts)
-    )
-
-    # rm directories
-    import shutil
-
-    shutil.rmtree("tests/test_solver/tmp")
+    )