diff --git a/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/large_file/large_file/large_file_mongo.py index 0376ce9..6e9d46e 100644 --- a/great_ai/large_file/large_file/large_file_mongo.py +++ b/great_ai/large_file/large_file/large_file_mongo.py @@ -65,7 +65,9 @@ class LargeFileMongo(LargeFileBase): for f in self._client.find( { "filename": re.compile( - re.escape(self._name + MONGO_NAME_VERSION_SEPARATOR) + ".*" + "^" + + re.escape(self._name + MONGO_NAME_VERSION_SEPARATOR) + + r"[0-9]+$" ) } ) diff --git a/tests/large_file/test_storage_safety.py b/tests/large_file/test_storage_safety.py index b4907d3..6b9a083 100644 --- a/tests/large_file/test_storage_safety.py +++ b/tests/large_file/test_storage_safety.py @@ -1,11 +1,17 @@ import io import tarfile from pathlib import Path +from types import SimpleNamespace from typing import Any, Iterator from unittest.mock import MagicMock, patch import pytest -from great_ai.large_file import LargeFileBase, LargeFileLocal, LargeFileS3 +from great_ai.large_file import ( + LargeFileBase, + LargeFileLocal, + LargeFileMongo, + LargeFileS3, +) @pytest.fixture(autouse=True) @@ -100,3 +106,18 @@ def test_evicted_remote_model_is_downloaded_again() -> None: cached.unlink() assert reader.get().read_text() == "remote weights" assert downloads == ["model/0", "model/0"] + + +def test_gridfs_matches_only_exact_model_names() -> None: + files = [ + SimpleNamespace(name=name, _id=name, length=1) + for name in ("model_0", "model_other_9", "other_model_10", "model_1_backup") + ] + backend = MagicMock() + backend.find.side_effect = lambda query: [ + file for file in files if query["filename"].search(file.name) + ] + with patch.object(LargeFileMongo, "_client", backend): + reader = LargeFileMongo("model") + assert reader.version == 0 + assert [instance.name for instance in reader._instances] == ["model"]