diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py index 92fa5d9..696b8a3 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -85,7 +85,7 @@ class LargeFileBase(ABC): self._newline = newline - LargeFileBase.cache_path.mkdir(parents=True, exist_ok=True) + self.cache_path.mkdir(parents=True, exist_ok=True) self._find_instances() self._check_mode_and_set_version() diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py new file mode 100644 index 0000000..89f9750 --- /dev/null +++ b/tests/large_file/test_storage_safety.py @@ -0,0 +1,22 @@ +from pathlib import Path +from typing import Iterator + +import pytest +from great_ai.large_file import LargeFileBase, LargeFileLocal + + +@pytest.fixture(autouse=True) +def storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + for implementation in (LargeFileBase, LargeFileLocal): + monkeypatch.setattr(implementation, "cache_path", tmp_path / "cache") + monkeypatch.setattr(implementation, "max_cache_size", None) + yield + + +def test_configured_cache_directory_is_created( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(LargeFileLocal, "cache_path", tmp_path / "custom-cache") + with LargeFileLocal("model", "w") as file: + file.write("weights") + assert LargeFileLocal("model").get().read_text() == "weights"