make file cleanup unified to fix assert bug in line 117

xopt-org · Dec 19, 2023 · 1ebcbf2 · 1ebcbf2
1 parent 41c35aa
commit 1ebcbf2
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/tests/generators/bayesian/test_high_level.py b/tests/generators/bayesian/test_high_level.py
@@ -132,8 +132,6 @@ def test_restart_torch_inline_serialization(self):
 
         X2.step()
 
-        os.remove("dump_inline.yml")
-
     def test_restart_torch_serialization(self):
         YAML = """
                 dump_file: dump.yml
@@ -177,9 +175,6 @@ def test_restart_torch_serialization(self):
 
         X2.step()
 
-        os.remove("generator_model.pt")
-        os.remove("dump.yml")
-
     def test_restart(self):
         YAML = """
                 dump_file: dump.yml
@@ -205,7 +200,6 @@ def test_restart(self):
         X.step()
 
         config = yaml.safe_load(open("dump.yml"))
-        os.remove("dump.yml")
 
         # test restart
         X2 = Xopt.model_validate(config)
@@ -222,7 +216,7 @@ def test_restart(self):
     @pytest.fixture(scope="module", autouse=True)
     def clean_up(self):
         yield
-        files = ["dump.yml", "mobo_model.pt", "dump_inline.yml"]
+        files = ["dump.yml", "mobo_model.pt", "dump_inline.yml", "generator_model.pt"]
         for f in files:
             if os.path.exists(f):
                 os.remove(f)