clean up tests that make files

ur-whitelab · Feb 23, 2024 · ef0b6b9 · ef0b6b9
1 parent 38e8cf1
commit ef0b6b9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/mdagent/tools/base_tools/analysis_tools/vis_tools.py b/mdagent/tools/base_tools/analysis_tools/vis_tools.py
@@ -69,10 +69,10 @@ def create_notebook(self, cif_file: str) -> str:
 
         # Code to import NGLview and display a file
         import_code = f"""
-        import nglview as nv
-        view = nv.show_file("{cif_file}")
-        view
-        """
+import nglview as nv
+view = nv.show_file("{cif_file}")
+view
+"""
 
         # Create new code cells
         install_cell = nbf.v4.new_code_cell(source=install_code)

diff --git a/tests/test_fxns.py b/tests/test_fxns.py
@@ -161,11 +161,15 @@ def test_find_png(vis_fxns):
 
 def test_create_notebook(path_to_cif, vis_fxns):
     result = vis_fxns.create_notebook(path_to_cif)
+    path_to_notebook = path_to_cif.split(".")[0] + "_vis.ipynb"
+    os.remove(path_to_notebook)
     assert result == "Visualization Complete"
 
 
 def test_add_hydrogens_and_remove_water(path_to_cif, cleaning_fxns):
     result = cleaning_fxns._add_hydrogens_and_remove_water(path_to_cif)
+    path_to_cleaned_file = "tidy_" + path_to_cif
+    os.remove(path_to_cleaned_file)
     assert "Cleaned File" in result