|
| 1 | +from collections.abc import Iterator |
| 2 | +from unittest.mock import MagicMock, patch |
| 3 | +from urllib.parse import parse_qs, urlsplit |
| 4 | + |
| 5 | +import pytest |
| 6 | +from google.api_core.exceptions import Forbidden, NotFound, ServiceUnavailable |
| 7 | +from google.auth.credentials import AnonymousCredentials |
| 8 | +from google.cloud.storage import Blob, Bucket, Client |
| 9 | +from google.oauth2.credentials import Credentials |
| 10 | +from requests import Response |
| 11 | + |
| 12 | +from tilebox.workflows.cache import GoogleStorageCache |
| 13 | + |
| 14 | + |
| 15 | +@pytest.fixture(params=[False, True], ids=["authenticated", "anonymous"]) |
| 16 | +def google_bucket(request: pytest.FixtureRequest) -> Iterator[Bucket]: |
| 17 | + credentials = AnonymousCredentials() if request.param else Credentials(token="configured-token") # noqa: S106 |
| 18 | + client = Client( |
| 19 | + project="configured-project", |
| 20 | + credentials=credentials, |
| 21 | + _http=MagicMock(), |
| 22 | + client_options={"api_endpoint": "https://custom.example.test"}, |
| 23 | + ) |
| 24 | + with ( |
| 25 | + patch("google.auth.default", side_effect=AssertionError("Bucket credentials must not be replaced by ADC")), |
| 26 | + patch("obstore.store.GCSStore", side_effect=AssertionError("Bucket configuration must stay with its client")), |
| 27 | + ): |
| 28 | + yield client.bucket("configured-bucket", user_project="billing-project") |
| 29 | + |
| 30 | + |
| 31 | +def test_google_bucket_read_preserves_client(google_bucket: Bucket) -> None: |
| 32 | + # Bucket reads and nested groups retain credentials, custom transport, endpoint, and requester-pays billing. |
| 33 | + credentials = google_bucket.client._credentials |
| 34 | + transport = google_bucket.client._http |
| 35 | + response = Response() |
| 36 | + response.status_code = 200 |
| 37 | + response._content = b"configured client content" |
| 38 | + response._content_consumed = True |
| 39 | + response.raw = MagicMock(headers={}) |
| 40 | + response.headers["content-length"] = str(len(response.content)) |
| 41 | + transport.request.return_value = response |
| 42 | + root = GoogleStorageCache(google_bucket) |
| 43 | + cache = root.group("folder").group("nested") |
| 44 | + assert isinstance(cache, GoogleStorageCache) |
| 45 | + assert cache.bucket is root.bucket is google_bucket |
| 46 | + assert cache["file.txt"] == b"configured client content" |
| 47 | + assert google_bucket.client._credentials is credentials |
| 48 | + # Newer Google clients may also fetch bucket metadata in a background thread. |
| 49 | + downloads = [call for call in transport.request.call_args_list if call.args and "/download/" in call.args[1]] |
| 50 | + assert len(downloads) == 1 |
| 51 | + method, url = downloads[0].args |
| 52 | + assert method == "GET" |
| 53 | + parsed = urlsplit(url) |
| 54 | + assert parsed.netloc == "custom.example.test" |
| 55 | + assert parsed.path.endswith("/configured-bucket/o/jobs%2Ffolder%2Fnested%2Ffile.txt") |
| 56 | + assert parse_qs(parsed.query)["userProject"] == ["billing-project"] |
| 57 | + |
| 58 | + |
| 59 | +@pytest.mark.parametrize("prefix", ["jobs", ""]) |
| 60 | +def test_google_bucket_operations(google_bucket: Bucket, prefix: str) -> None: |
| 61 | + # Bucket caches preserve writes, membership, deletion, and delimiter-based listing for each group. |
| 62 | + cache = GoogleStorageCache(google_bucket, prefix).group("folder") |
| 63 | + path = f"{prefix}/folder" if prefix else "folder" |
| 64 | + with ( |
| 65 | + patch.object(Blob, "upload_from_file", autospec=True) as upload, |
| 66 | + patch.object(Blob, "exists", autospec=True, return_value=True) as exists, |
| 67 | + patch.object(Blob, "delete", autospec=True) as delete, |
| 68 | + patch.object(google_bucket, "list_blobs", return_value=[Blob(f"{path}/one", google_bucket)]) as listing, |
| 69 | + ): |
| 70 | + cache["one"] = b"saved value" |
| 71 | + blob, stream = upload.call_args.args |
| 72 | + assert blob.bucket is google_bucket |
| 73 | + assert blob.name == f"{path}/one" |
| 74 | + assert stream.getvalue() == b"saved value" |
| 75 | + assert "one" in cache |
| 76 | + assert exists.call_args.args[0].name == f"{path}/one" |
| 77 | + assert list(cache) == ["one"] |
| 78 | + listing.assert_called_once_with(prefix=f"{path}/", delimiter="/") |
| 79 | + del cache["one"] |
| 80 | + assert delete.call_args.args[0].name == f"{path}/one" |
| 81 | + |
| 82 | + |
| 83 | +@pytest.mark.parametrize("error_type", [NotFound, Forbidden, ServiceUnavailable]) |
| 84 | +def test_google_bucket_read_errors(google_bucket: Bucket, error_type: type[Exception]) -> None: |
| 85 | + # Only missing objects become KeyError; configured-client permission and backend errors propagate unchanged. |
| 86 | + error = error_type("read failed") |
| 87 | + with patch.object(Blob, "download_as_bytes", side_effect=error): |
| 88 | + expected = KeyError if error_type is NotFound else error_type |
| 89 | + with pytest.raises(expected) as raised: |
| 90 | + GoogleStorageCache(google_bucket).group("nested")["file"] |
| 91 | + if error_type is not NotFound: |
| 92 | + assert raised.value is error |
0 commit comments