function to clear ckpts

mdagent/utils/clear_mem.py

@@ -0,0 +1,31 @@
+import glob
+import os
+import shutil
+
+
+class ClearMemory:
+    def _clear_ckpts(self):
+        directories_to_remove = ["files/pdb", "files/simulations", "files/records"]
+
+        for directory in directories_to_remove:
+            if os.path.exists(directory):
+                shutil.rmtree(directory)
+
+        temp_files = glob.glob("temp_*")
+        for temp_file in temp_files:
+            if os.path.isfile(temp_file):
+                os.remove(temp_file)
+
+        if os.path.exists("paths_registry.json"):
+            os.remove("paths_registry.json")
+
+    def clear_ckpts_root(self):
+        current_dir = os.getcwd()
+        while current_dir != os.path.dirname(current_dir):
+            if "setup.py" in os.listdir(current_dir):
+                os.chdir(current_dir)
+                self._clear_ckpts()
+                break
+            else:
+                current_dir = os.path.dirname(current_dir)
+        os.chdir(current_dir)  # Reset the working directory

tests/test_utils.py

@@ -0,0 +1,52 @@
+import os
+
+import pytest
+
+from mdagent.utils.clear_mem import ClearMemory
+
+
+@pytest.fixture
+def clear_memory():
+    return ClearMemory()
+
+
+@pytest.fixture
+def setup_test_environment(tmp_path):
+    # temporary test environment
+    base_dir = tmp_path / "test_env"
+    base_dir.mkdir()
+    (base_dir / "files" / "pdb").mkdir(parents=True)
+    (base_dir / "files" / "simulations").mkdir(parents=True)
+    (base_dir / "files" / "records").mkdir(parents=True)
+    (base_dir / "temp_file").write_text("Temporary file content")
+    (base_dir / "paths_registry.json").write_text("{}")
+    return base_dir
+
+
+def test_clear_ckpts(setup_test_environment, clear_memory):
+    base_dir = setup_test_environment
+    os.chdir(base_dir)
+    clear_memory._clear_ckpts()
+
+    # Check that the directories and files have been removed
+    assert not os.path.exists(base_dir / "files" / "pdb")
+    assert not os.path.exists(base_dir / "files" / "simulations")
+    assert not os.path.exists(base_dir / "files" / "records")
+    assert not os.path.exists(base_dir / "temp_file")
+    assert not os.path.exists(base_dir / "paths_registry.json")
+
+
+def test_clear_ckpts_root(setup_test_environment, clear_memory):
+    base_dir = setup_test_environment
+    os.chdir(base_dir / "files")
+    (base_dir / "setup.py").write_text("# setup.py content")
+
+    clear_memory.clear_ckpts_root()
+    assert os.getcwd() == str(base_dir)
+
+    # Check that the directories and files have been removed
+    assert not os.path.exists(base_dir / "files" / "pdb")
+    assert not os.path.exists(base_dir / "files" / "simulations")
+    assert not os.path.exists(base_dir / "files" / "records")
+    assert not os.path.exists(base_dir / "temp_file")
+    assert not os.path.exists(base_dir / "paths_registry.json")