diff --git a/mdagent/utils/clear_mem.py b/mdagent/utils/clear_mem.py new file mode 100644 index 00000000..92c743db --- /dev/null +++ b/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 diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..04c2886c --- /dev/null +++ b/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")