wehub-resource-sync
ec2b666284
Continuous Integration / Pre-commit Linter (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.10) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.11) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.12) (push) Has been cancelled
Continuous Integration / Mypy Check (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.12) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / Unit Tests (Python 3.14) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.10) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.11) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.12) (push) Has been cancelled
Copybara PR Handler / close-imported-pr (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.13) (push) Has been cancelled
Continuous Integration / A2A v0.3 Tests (Python 3.14) (push) Has been cancelled
1618 lines
45 KiB
Python
1618 lines
45 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# pylint: disable=missing-class-docstring,missing-function-docstring
|
|
|
|
"""Tests for the artifact service."""
|
|
|
|
from datetime import datetime
|
|
import enum
|
|
import json
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from typing import Optional
|
|
from typing import Union
|
|
from unittest import mock
|
|
from unittest.mock import patch
|
|
from urllib.parse import unquote
|
|
from urllib.parse import urlparse
|
|
|
|
from google.adk.artifacts import file_artifact_service
|
|
from google.adk.artifacts.base_artifact_service import ArtifactVersion
|
|
from google.adk.artifacts.base_artifact_service import ensure_part
|
|
from google.adk.artifacts.file_artifact_service import FileArtifactService
|
|
from google.adk.artifacts.gcs_artifact_service import GcsArtifactService
|
|
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
|
|
from google.adk.errors.input_validation_error import InputValidationError
|
|
from google.genai import types
|
|
import pytest
|
|
|
|
Enum = enum.Enum
|
|
|
|
# Define a fixed datetime object to be returned by datetime.now()
|
|
FIXED_DATETIME = datetime(2025, 1, 1, 12, 0, 0)
|
|
|
|
|
|
class ArtifactServiceType(Enum):
|
|
FILE = "FILE"
|
|
IN_MEMORY = "IN_MEMORY"
|
|
GCS = "GCS"
|
|
|
|
|
|
class MockBlob:
|
|
"""Mocks a GCS Blob object.
|
|
|
|
This class provides mock implementations for a few common GCS Blob methods,
|
|
allowing the user to test code that interacts with GCS without actually
|
|
connecting to a real bucket.
|
|
"""
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Initializes a MockBlob.
|
|
|
|
Args:
|
|
name: The name of the blob.
|
|
"""
|
|
self.name = name
|
|
self.content: Optional[bytes] = None
|
|
self.content_type: Optional[str] = None
|
|
self.time_created = FIXED_DATETIME
|
|
self.metadata: dict[str, Any] = {}
|
|
|
|
def upload_from_string(
|
|
self, data: Union[str, bytes], content_type: Optional[str] = None
|
|
) -> None:
|
|
"""Mocks uploading data to the blob (from a string or bytes).
|
|
|
|
Args:
|
|
data: The data to upload (string or bytes).
|
|
content_type: The content type of the data (optional).
|
|
"""
|
|
if isinstance(data, str):
|
|
self.content = data.encode("utf-8")
|
|
elif isinstance(data, bytes):
|
|
self.content = data
|
|
else:
|
|
raise TypeError("data must be str or bytes")
|
|
|
|
if content_type:
|
|
self.content_type = content_type
|
|
|
|
def download_as_bytes(self) -> bytes:
|
|
"""Mocks downloading the blob's content as bytes.
|
|
|
|
Returns:
|
|
bytes: The content of the blob as bytes.
|
|
|
|
Raises:
|
|
Exception: If the blob doesn't exist (hasn't been uploaded to).
|
|
"""
|
|
if self.content is None:
|
|
return b""
|
|
return self.content
|
|
|
|
def delete(self) -> None:
|
|
"""Mocks deleting a blob."""
|
|
self.content = None
|
|
self.content_type = None
|
|
|
|
|
|
class MockBucket:
|
|
"""Mocks a GCS Bucket object."""
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Initializes a MockBucket.
|
|
|
|
Args:
|
|
name: The name of the bucket.
|
|
"""
|
|
self.name = name
|
|
self.blobs: dict[str, MockBlob] = {}
|
|
|
|
def blob(self, blob_name: str) -> MockBlob:
|
|
"""Mocks getting a Blob object (doesn't create it in storage).
|
|
|
|
Args:
|
|
blob_name: The name of the blob.
|
|
|
|
Returns:
|
|
A MockBlob instance.
|
|
"""
|
|
if blob_name not in self.blobs:
|
|
self.blobs[blob_name] = MockBlob(blob_name)
|
|
return self.blobs[blob_name]
|
|
|
|
def get_blob(self, blob_name: str) -> Optional[MockBlob]:
|
|
"""Mocks getting a blob from storage if it exists and has content."""
|
|
blob = self.blobs.get(blob_name)
|
|
if blob and blob.content is not None:
|
|
return blob
|
|
return None
|
|
|
|
|
|
class MockClient:
|
|
"""Mocks the GCS Client."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Initializes MockClient."""
|
|
self.buckets: dict[str, MockBucket] = {}
|
|
|
|
def bucket(self, bucket_name: str) -> MockBucket:
|
|
"""Mocks getting a Bucket object."""
|
|
if bucket_name not in self.buckets:
|
|
self.buckets[bucket_name] = MockBucket(bucket_name)
|
|
return self.buckets[bucket_name]
|
|
|
|
def list_blobs(self, bucket: MockBucket, prefix: Optional[str] = None):
|
|
"""Mocks listing blobs in a bucket, optionally with a prefix."""
|
|
if prefix:
|
|
return [
|
|
blob
|
|
for name, blob in bucket.blobs.items()
|
|
if name.startswith(prefix) and blob.content is not None
|
|
]
|
|
return [blob for blob in bucket.blobs.values() if blob.content is not None]
|
|
|
|
|
|
def mock_gcs_artifact_service():
|
|
with mock.patch("google.cloud.storage.Client", return_value=MockClient()):
|
|
return GcsArtifactService(bucket_name="test_bucket")
|
|
|
|
|
|
@pytest.fixture
|
|
def artifact_service_factory(tmp_path: Path):
|
|
"""Provides an artifact service constructor bound to the test tmp path."""
|
|
|
|
def factory(
|
|
service_type: ArtifactServiceType = ArtifactServiceType.IN_MEMORY,
|
|
):
|
|
if service_type == ArtifactServiceType.GCS:
|
|
return mock_gcs_artifact_service()
|
|
if service_type == ArtifactServiceType.FILE:
|
|
return FileArtifactService(root_dir=tmp_path / "artifacts")
|
|
return InMemoryArtifactService()
|
|
|
|
return factory
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_load_empty(service_type, artifact_service_factory):
|
|
"""Tests loading an artifact when none exists."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
assert not await artifact_service.load_artifact(
|
|
app_name="test_app",
|
|
user_id="test_user",
|
|
session_id="session_id",
|
|
filename="filename",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_save_load_delete(service_type, artifact_service_factory):
|
|
"""Tests saving, loading, and deleting an artifact."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
filename = "file456"
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=artifact,
|
|
)
|
|
assert (
|
|
await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
== artifact
|
|
)
|
|
|
|
# Attempt to load a version that doesn't exist
|
|
assert not await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
version=3,
|
|
)
|
|
|
|
await artifact_service.delete_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
assert not await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_in_memory_loads_nested_artifact_reference(
|
|
artifact_service_factory,
|
|
):
|
|
"""Tests loading an artifact reference whose target name is nested."""
|
|
artifact_service = artifact_service_factory(ArtifactServiceType.IN_MEMORY)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
target_filename = "folder/file456"
|
|
target_artifact = types.Part.from_text(text="target")
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=target_filename,
|
|
artifact=target_artifact,
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename="reference",
|
|
artifact=types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/app0/users/user0/sessions/123/artifacts/"
|
|
"folder/file456/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
),
|
|
)
|
|
|
|
assert (
|
|
await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename="reference",
|
|
)
|
|
== target_artifact
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_list_keys(service_type, artifact_service_factory):
|
|
"""Tests listing keys in the artifact service."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
filename = "filename"
|
|
filenames = [filename + str(i) for i in range(5)]
|
|
|
|
for f in filenames:
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=f,
|
|
artifact=artifact,
|
|
)
|
|
|
|
assert (
|
|
await artifact_service.list_artifact_keys(
|
|
app_name=app_name, user_id=user_id, session_id=session_id
|
|
)
|
|
== filenames
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_list_versions(service_type, artifact_service_factory):
|
|
"""Tests listing versions of an artifact."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
filename = "with/slash/filename"
|
|
versions = [
|
|
types.Part.from_bytes(
|
|
data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
|
|
)
|
|
for i in range(3)
|
|
]
|
|
versions.append(types.Part.from_text(text="hello"))
|
|
|
|
for i in range(4):
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=versions[i],
|
|
)
|
|
|
|
response_versions = await artifact_service.list_versions(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
|
|
assert response_versions == list(range(4))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_list_keys_preserves_user_prefix(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests that list_artifact_keys preserves 'user:' prefix in returned names."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
|
|
# Save artifacts with "user:" prefix (cross-session artifacts)
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename="user:document.pdf",
|
|
artifact=artifact,
|
|
)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename="user:image.png",
|
|
artifact=artifact,
|
|
)
|
|
|
|
# Save session-scoped artifact without prefix
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename="session_file.txt",
|
|
artifact=artifact,
|
|
)
|
|
|
|
# List artifacts should return names with "user:" prefix for user-scoped artifacts
|
|
artifact_keys = await artifact_service.list_artifact_keys(
|
|
app_name=app_name, user_id=user_id, session_id=session_id
|
|
)
|
|
|
|
# Should contain prefixed names and session file
|
|
expected_keys = ["user:document.pdf", "user:image.png", "session_file.txt"]
|
|
assert sorted(artifact_keys) == sorted(expected_keys)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
|
|
)
|
|
async def test_list_artifact_versions_and_get_artifact_version(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests listing artifact versions and getting a specific version."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
filename = "filename"
|
|
versions = [
|
|
types.Part.from_bytes(
|
|
data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
|
|
)
|
|
for i in range(4)
|
|
]
|
|
|
|
with patch(
|
|
"google.adk.artifacts.base_artifact_service.platform_time"
|
|
) as mock_platform_time:
|
|
mock_platform_time.get_time.return_value = FIXED_DATETIME.timestamp()
|
|
|
|
for i in range(4):
|
|
custom_metadata = {"key": "value" + str(i)}
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=versions[i],
|
|
custom_metadata=custom_metadata,
|
|
)
|
|
|
|
artifact_versions = await artifact_service.list_artifact_versions(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
|
|
expected_artifact_versions = []
|
|
for i in range(4):
|
|
metadata = {"key": "value" + str(i)}
|
|
if service_type == ArtifactServiceType.GCS:
|
|
uri = (
|
|
f"gs://test_bucket/{app_name}/{user_id}/{session_id}/{filename}/{i}"
|
|
)
|
|
else:
|
|
uri = f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{i}"
|
|
expected_artifact_versions.append(
|
|
ArtifactVersion(
|
|
version=i,
|
|
canonical_uri=uri,
|
|
custom_metadata=metadata,
|
|
mime_type="text/plain",
|
|
create_time=FIXED_DATETIME.timestamp(),
|
|
)
|
|
)
|
|
assert artifact_versions == expected_artifact_versions
|
|
|
|
# Get latest artifact version when version is not specified
|
|
assert (
|
|
await artifact_service.get_artifact_version(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
== expected_artifact_versions[-1]
|
|
)
|
|
|
|
# Get artifact version by version number
|
|
assert (
|
|
await artifact_service.get_artifact_version(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
version=2,
|
|
)
|
|
== expected_artifact_versions[2]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
|
|
)
|
|
async def test_list_artifact_versions_with_user_prefix(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests listing artifact versions with user prefix."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
user_scoped_filename = "user:document.pdf"
|
|
versions = [
|
|
types.Part.from_bytes(
|
|
data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
|
|
)
|
|
for i in range(4)
|
|
]
|
|
|
|
with patch(
|
|
"google.adk.artifacts.base_artifact_service.platform_time"
|
|
) as mock_platform_time:
|
|
mock_platform_time.get_time.return_value = FIXED_DATETIME.timestamp()
|
|
|
|
for i in range(4):
|
|
custom_metadata = {"key": "value" + str(i)}
|
|
# Save artifacts with "user:" prefix (cross-session artifacts)
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=user_scoped_filename,
|
|
artifact=versions[i],
|
|
custom_metadata=custom_metadata,
|
|
)
|
|
|
|
artifact_versions = await artifact_service.list_artifact_versions(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=user_scoped_filename,
|
|
)
|
|
|
|
expected_artifact_versions = []
|
|
for i in range(4):
|
|
metadata = {"key": "value" + str(i)}
|
|
if service_type == ArtifactServiceType.GCS:
|
|
uri = f"gs://test_bucket/{app_name}/{user_id}/user/{user_scoped_filename}/{i}"
|
|
else:
|
|
uri = f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}"
|
|
expected_artifact_versions.append(
|
|
ArtifactVersion(
|
|
version=i,
|
|
canonical_uri=uri,
|
|
custom_metadata=metadata,
|
|
mime_type="text/plain",
|
|
create_time=FIXED_DATETIME.timestamp(),
|
|
)
|
|
)
|
|
assert artifact_versions == expected_artifact_versions
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
|
|
)
|
|
async def test_get_artifact_version_artifact_does_not_exist(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests getting an artifact version when artifact does not exist."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
assert not await artifact_service.get_artifact_version(
|
|
app_name="test_app",
|
|
user_id="test_user",
|
|
session_id="session_id",
|
|
filename="filename",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
|
|
)
|
|
async def test_get_artifact_version_out_of_index(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests loading an artifact with an out-of-index version."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "123"
|
|
filename = "filename"
|
|
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=artifact,
|
|
)
|
|
|
|
# Attempt to get a version that doesn't exist
|
|
assert not await artifact_service.get_artifact_version(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
version=3,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gcs_save_and_load_empty_text_artifact(
|
|
artifact_service_factory,
|
|
):
|
|
"""GcsArtifactService should round-trip empty text as text."""
|
|
artifact_service = artifact_service_factory(ArtifactServiceType.GCS)
|
|
artifact = types.Part.from_text(text="")
|
|
|
|
version = await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="empty.txt",
|
|
artifact=artifact,
|
|
)
|
|
|
|
assert version == 0
|
|
loaded_artifact = await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="empty.txt",
|
|
)
|
|
|
|
assert loaded_artifact == types.Part(text="")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_metadata_camelcase(tmp_path, artifact_service_factory):
|
|
"""Ensures FileArtifactService writes camelCase metadata without newlines."""
|
|
artifact_service = artifact_service_factory(ArtifactServiceType.FILE)
|
|
artifact = types.Part.from_bytes(
|
|
data=b"binary-content", mime_type="application/octet-stream"
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id="sess789",
|
|
filename="docs/report.txt",
|
|
artifact=artifact,
|
|
)
|
|
|
|
metadata_path = (
|
|
tmp_path
|
|
/ "artifacts"
|
|
/ "users"
|
|
/ "user123"
|
|
/ "sessions"
|
|
/ "sess789"
|
|
/ "artifacts"
|
|
/ "docs"
|
|
/ "report.txt"
|
|
/ "versions"
|
|
/ "0"
|
|
/ "metadata.json"
|
|
)
|
|
raw_metadata = metadata_path.read_text(encoding="utf-8")
|
|
assert "\n" not in raw_metadata
|
|
|
|
metadata = json.loads(raw_metadata)
|
|
payload_path = (metadata_path.parent / "report.txt").resolve()
|
|
expected_canonical_uri = payload_path.as_uri()
|
|
create_time = metadata.pop("createTime", None)
|
|
assert create_time is not None
|
|
assert metadata == {
|
|
"fileName": "docs/report.txt",
|
|
"mimeType": "application/octet-stream",
|
|
"canonicalUri": expected_canonical_uri,
|
|
"version": 0,
|
|
"customMetadata": {},
|
|
}
|
|
parsed_canonical = urlparse(metadata["canonicalUri"])
|
|
canonical_path = Path(unquote(parsed_canonical.path))
|
|
assert canonical_path.name == "report.txt"
|
|
assert canonical_path.read_bytes() == b"binary-content"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_list_artifact_versions(tmp_path, artifact_service_factory):
|
|
"""FileArtifactService exposes canonical URIs and metadata for each version."""
|
|
artifact_service = artifact_service_factory(ArtifactServiceType.FILE)
|
|
artifact = types.Part.from_bytes(
|
|
data=b"binary-content", mime_type="application/octet-stream"
|
|
)
|
|
custom_metadata = {"origin": "unit-test"}
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id="sess789",
|
|
filename="docs/report.txt",
|
|
artifact=artifact,
|
|
custom_metadata=custom_metadata,
|
|
)
|
|
|
|
versions = await artifact_service.list_artifact_versions(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id="sess789",
|
|
filename="docs/report.txt",
|
|
)
|
|
assert len(versions) == 1
|
|
version_meta = versions[0]
|
|
assert version_meta.version == 0
|
|
version_payload_path = (
|
|
tmp_path
|
|
/ "artifacts"
|
|
/ "users"
|
|
/ "user123"
|
|
/ "sessions"
|
|
/ "sess789"
|
|
/ "artifacts"
|
|
/ "docs"
|
|
/ "report.txt"
|
|
/ "versions"
|
|
/ "0"
|
|
/ "report.txt"
|
|
).resolve()
|
|
assert version_meta.canonical_uri == version_payload_path.as_uri()
|
|
assert version_meta.custom_metadata == custom_metadata
|
|
parsed_version_uri = urlparse(version_meta.canonical_uri)
|
|
version_uri_path = Path(unquote(parsed_version_uri.path))
|
|
assert version_uri_path.read_bytes() == b"binary-content"
|
|
|
|
fetched = await artifact_service.get_artifact_version(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id="sess789",
|
|
filename="docs/report.txt",
|
|
version=0,
|
|
)
|
|
assert fetched is not None
|
|
assert fetched.version == version_meta.version
|
|
assert fetched.canonical_uri == version_meta.canonical_uri
|
|
assert fetched.custom_metadata == version_meta.custom_metadata
|
|
|
|
latest = await artifact_service.get_artifact_version(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id="sess789",
|
|
filename="docs/report.txt",
|
|
)
|
|
assert latest is not None
|
|
assert latest.version == version_meta.version
|
|
assert latest.canonical_uri == version_meta.canonical_uri
|
|
assert latest.custom_metadata == version_meta.custom_metadata
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("filename", "session_id"),
|
|
[
|
|
("../escape.txt", "sess123"),
|
|
("user:../escape.txt", "sess123"),
|
|
("/absolute/path.txt", "sess123"),
|
|
("user:/absolute/path.txt", None),
|
|
],
|
|
)
|
|
async def test_file_save_artifact_rejects_out_of_scope_paths(
|
|
tmp_path, filename, session_id
|
|
):
|
|
"""FileArtifactService prevents path traversal outside of its storage roots."""
|
|
artifact_service = FileArtifactService(root_dir=tmp_path / "artifacts")
|
|
part = types.Part(text="content")
|
|
with pytest.raises(InputValidationError):
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=part,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"user_id",
|
|
[
|
|
"../escape",
|
|
"../../etc",
|
|
"foo/../../bar",
|
|
"valid/../..",
|
|
"..",
|
|
".",
|
|
"has/slash",
|
|
"back\\slash",
|
|
"null\x00byte",
|
|
"",
|
|
],
|
|
)
|
|
async def test_file_save_artifact_rejects_traversal_in_user_id(
|
|
tmp_path, user_id
|
|
):
|
|
"""FileArtifactService rejects user_id values that escape root_dir."""
|
|
artifact_service = FileArtifactService(root_dir=tmp_path / "artifacts")
|
|
part = types.Part(text="content")
|
|
with pytest.raises(InputValidationError):
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id=user_id,
|
|
session_id="sess123",
|
|
filename="safe.txt",
|
|
artifact=part,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"session_id",
|
|
[
|
|
"../escape",
|
|
"../../tmp",
|
|
"foo/../../bar",
|
|
"..",
|
|
".",
|
|
"has/slash",
|
|
"back\\slash",
|
|
"null\x00byte",
|
|
"",
|
|
],
|
|
)
|
|
async def test_file_save_artifact_rejects_traversal_in_session_id(
|
|
tmp_path, session_id
|
|
):
|
|
"""FileArtifactService rejects session_id values that escape root_dir."""
|
|
artifact_service = FileArtifactService(root_dir=tmp_path / "artifacts")
|
|
part = types.Part(text="content")
|
|
with pytest.raises(InputValidationError):
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id=session_id,
|
|
filename="safe.txt",
|
|
artifact=part,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path):
|
|
"""Absolute filenames are rejected even when they point inside the scope."""
|
|
artifact_service = FileArtifactService(root_dir=tmp_path / "artifacts")
|
|
absolute_in_scope = (
|
|
tmp_path
|
|
/ "artifacts"
|
|
/ "apps"
|
|
/ "myapp"
|
|
/ "users"
|
|
/ "user123"
|
|
/ "artifacts"
|
|
/ "diagram.png"
|
|
)
|
|
part = types.Part(text="content")
|
|
with pytest.raises(InputValidationError):
|
|
await artifact_service.save_artifact(
|
|
app_name="myapp",
|
|
user_id="user123",
|
|
session_id=None,
|
|
filename=str(absolute_in_scope),
|
|
artifact=part,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
],
|
|
)
|
|
async def test_artifact_reference_allows_same_session_scope(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""ArtifactService allows references inside the same session scope."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="source.txt",
|
|
artifact=types.Part(text="hello"),
|
|
)
|
|
|
|
ref = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/app0/users/user0/sessions/sess0/"
|
|
"artifacts/source.txt/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="ref.txt",
|
|
artifact=ref,
|
|
)
|
|
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="ref.txt",
|
|
)
|
|
assert loaded == types.Part(text="hello")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
],
|
|
)
|
|
async def test_artifact_reference_allows_same_user_user_scope(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""ArtifactService allows references to user-scoped files from same user."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="user:profile.txt",
|
|
artifact=types.Part(text="profile"),
|
|
)
|
|
|
|
ref = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/app0/users/user0/artifacts/"
|
|
"user:profile.txt/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess1",
|
|
filename="ref.txt",
|
|
artifact=ref,
|
|
)
|
|
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess1",
|
|
filename="ref.txt",
|
|
)
|
|
assert loaded == types.Part(text="profile")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
],
|
|
)
|
|
async def test_artifact_reference_rejects_cross_user_on_save(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""ArtifactService rejects references to different users on save."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="victim",
|
|
session_id="victim-sess",
|
|
filename="user:secret.txt",
|
|
artifact=types.Part(text="secret"),
|
|
)
|
|
|
|
ref = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/app0/users/victim/artifacts/"
|
|
"user:secret.txt/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
with pytest.raises(InputValidationError, match="same app and user scope"):
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="attacker",
|
|
session_id="attacker-sess",
|
|
filename="ref.txt",
|
|
artifact=ref,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
],
|
|
)
|
|
async def test_artifact_reference_rejects_cross_app_on_save(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""ArtifactService rejects references to different apps on save."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="victim-app",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="user:secret.txt",
|
|
artifact=types.Part(text="secret"),
|
|
)
|
|
|
|
ref = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/victim-app/users/user0/artifacts/"
|
|
"user:secret.txt/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
with pytest.raises(InputValidationError, match="same app and user scope"):
|
|
await artifact_service.save_artifact(
|
|
app_name="attacker-app",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="ref.txt",
|
|
artifact=ref,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
],
|
|
)
|
|
async def test_artifact_reference_rejects_cross_session_on_load(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""ArtifactService rejects modified references to different sessions on load."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="source.txt",
|
|
artifact=types.Part(text="source"),
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess1",
|
|
filename="source.txt",
|
|
artifact=types.Part(text="other-session"),
|
|
)
|
|
|
|
ref = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=(
|
|
"artifact://apps/app0/users/user0/sessions/sess0/"
|
|
"artifacts/source.txt/versions/0"
|
|
),
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="ref.txt",
|
|
artifact=ref,
|
|
)
|
|
|
|
new_uri = (
|
|
"artifact://apps/app0/users/user0/sessions/sess1/"
|
|
"artifacts/source.txt/versions/0"
|
|
)
|
|
# Manually modify the stored reference URI to point to a different session.
|
|
if service_type == ArtifactServiceType.GCS:
|
|
blob_name = artifact_service._get_blob_name(
|
|
"app0", "user0", "ref.txt", 0, "sess0"
|
|
)
|
|
blob = artifact_service.bucket.get_blob(blob_name)
|
|
blob.metadata["adkFileUri"] = new_uri
|
|
elif service_type == ArtifactServiceType.IN_MEMORY:
|
|
ref_path = artifact_service._artifact_path(
|
|
"app0", "user0", "ref.txt", "sess0"
|
|
)
|
|
artifact_service.artifacts[ref_path][0].data.file_data.file_uri = new_uri
|
|
|
|
with pytest.raises(InputValidationError, match="same session scope"):
|
|
await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="sess0",
|
|
filename="ref.txt",
|
|
)
|
|
|
|
|
|
class TestEnsurePart:
|
|
"""Tests for the ensure_part normalization helper."""
|
|
|
|
def test_returns_part_unchanged(self):
|
|
"""A types.Part instance passes through without modification."""
|
|
part = types.Part.from_bytes(data=b"hello", mime_type="text/plain")
|
|
result = ensure_part(part)
|
|
assert result is part
|
|
|
|
def test_converts_camel_case_dict(self):
|
|
"""A camelCase dict (Agentspace format) is converted to types.Part."""
|
|
raw = {"inlineData": {"mimeType": "image/png", "data": "dGVzdA=="}}
|
|
result = ensure_part(raw)
|
|
assert isinstance(result, types.Part)
|
|
assert result.inline_data is not None
|
|
assert result.inline_data.mime_type == "image/png"
|
|
|
|
def test_converts_snake_case_dict(self):
|
|
"""A snake_case dict is converted to types.Part."""
|
|
raw = {"inline_data": {"mime_type": "text/plain", "data": "aGVsbG8="}}
|
|
result = ensure_part(raw)
|
|
assert isinstance(result, types.Part)
|
|
assert result.inline_data is not None
|
|
assert result.inline_data.mime_type == "text/plain"
|
|
|
|
def test_converts_text_dict(self):
|
|
"""A dict with 'text' key is converted to types.Part."""
|
|
raw = {"text": "hello world"}
|
|
result = ensure_part(raw)
|
|
assert isinstance(result, types.Part)
|
|
assert result.text == "hello world"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GCS file_data (URI reference) tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_with_external_gcs_uri() -> None:
|
|
"""GcsArtifactService saves and loads a gs:// file_data URI reference."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
artifact = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri="gs://my-bucket/report.pdf",
|
|
mime_type="application/pdf",
|
|
)
|
|
)
|
|
|
|
version = await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="report.pdf",
|
|
artifact=artifact,
|
|
)
|
|
assert version == 0
|
|
|
|
loaded = await service.load_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="report.pdf",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.file_data is not None
|
|
assert loaded.file_data.file_uri == "gs://my-bucket/report.pdf"
|
|
assert loaded.file_data.mime_type == "application/pdf"
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_with_artifact_ref_uri() -> None:
|
|
"""GcsArtifactService saves and recursively loads an internal artifact:// URI reference."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
|
|
# Save the referenced (source) artifact first.
|
|
source_artifact = types.Part(text="source content")
|
|
await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="source.txt",
|
|
artifact=source_artifact,
|
|
)
|
|
|
|
artifact_ref_uri = "artifact://apps/app/users/user1/sessions/sess1/artifacts/source.txt/versions/0"
|
|
artifact = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri=artifact_ref_uri,
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
|
|
version = await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="ref.txt",
|
|
artifact=artifact,
|
|
)
|
|
assert version == 0
|
|
|
|
loaded = await service.load_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="ref.txt",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.text == "source content"
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_file_data_without_mime_type() -> None:
|
|
"""GcsArtifactService handles file_data with no mime_type."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
artifact = types.Part(
|
|
file_data=types.FileData(file_uri="gs://my-bucket/data.bin")
|
|
)
|
|
|
|
version = await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="data.bin",
|
|
artifact=artifact,
|
|
)
|
|
assert version == 0
|
|
|
|
loaded = await service.load_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="data.bin",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.file_data is not None
|
|
assert loaded.file_data.file_uri == "gs://my-bucket/data.bin"
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_file_data_missing_uri_raises() -> None:
|
|
"""GcsArtifactService raises InputValidationError when file_uri is empty."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
artifact = types.Part(file_data=types.FileData(file_uri=""))
|
|
|
|
with pytest.raises(InputValidationError):
|
|
await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="empty.bin",
|
|
artifact=artifact,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_file_data_invalid_uri_raises() -> None:
|
|
"""GcsArtifactService raises InputValidationError when file_uri is an invalid artifact:// URI template."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
artifact = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri="artifact://apps/app/invalid",
|
|
mime_type="text/plain",
|
|
)
|
|
)
|
|
|
|
with pytest.raises(InputValidationError):
|
|
await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="invalid_ref.txt",
|
|
artifact=artifact,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_save_artifact_metadata_namespacing_and_mime() -> None:
|
|
"""GcsArtifactService saves file_data using namespaced metadata keys."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
artifact = types.Part(
|
|
file_data=types.FileData(
|
|
file_uri="gs://my-bucket/report.pdf",
|
|
mime_type="application/pdf",
|
|
)
|
|
)
|
|
|
|
await service.save_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="report.pdf",
|
|
artifact=artifact,
|
|
)
|
|
|
|
blob_name = service._get_blob_name("app", "user1", "report.pdf", 0, "sess1")
|
|
blob = service.bucket.get_blob(blob_name)
|
|
assert blob is not None
|
|
assert blob.metadata.get("adkFileUri") == "gs://my-bucket/report.pdf"
|
|
assert blob.metadata.get("adkFileMimeType") == "application/pdf"
|
|
assert "file_uri" not in blob.metadata
|
|
|
|
|
|
@pytest.mark.asyncio # type: ignore[untyped-decorator]
|
|
async def test_gcs_load_artifact_file_data_fallback_compatibility() -> None:
|
|
"""GcsArtifactService loads file_data with old file_uri metadata key for backward compatibility."""
|
|
service = mock_gcs_artifact_service() # type: ignore[no-untyped-call]
|
|
blob_name = service._get_blob_name(
|
|
"app", "user1", "old_report.pdf", 0, "sess1"
|
|
)
|
|
blob = service.bucket.blob(blob_name)
|
|
# Manually setup metadata with old key
|
|
blob.metadata = {"file_uri": "gs://my-bucket/old_report.pdf"}
|
|
blob.upload_from_string(b"", content_type="application/pdf")
|
|
|
|
loaded = await service.load_artifact(
|
|
app_name="app",
|
|
user_id="user1",
|
|
session_id="sess1",
|
|
filename="old_report.pdf",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.file_data is not None
|
|
assert loaded.file_data.file_uri == "gs://my-bucket/old_report.pdf"
|
|
assert loaded.file_data.mime_type == "application/pdf"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_save_artifact_with_camel_case_dict(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Artifact services accept camelCase dicts (Agentspace format)."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "sess0"
|
|
filename = "uploaded.png"
|
|
|
|
# Simulate what Agentspace sends: a plain dict with camelCase keys.
|
|
raw_artifact = {
|
|
"inlineData": {
|
|
"mimeType": "image/png",
|
|
"data": "dGVzdF9pbWFnZV9kYXRh",
|
|
}
|
|
}
|
|
|
|
version = await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=raw_artifact,
|
|
)
|
|
assert version == 0
|
|
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.inline_data is not None
|
|
assert loaded.inline_data.mime_type == "image/png"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_save_artifact_with_snake_case_dict(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Artifact services accept snake_case dicts."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "sess0"
|
|
filename = "uploaded.txt"
|
|
|
|
raw_artifact = {
|
|
"inline_data": {
|
|
"mime_type": "text/plain",
|
|
"data": "aGVsbG8=",
|
|
}
|
|
}
|
|
|
|
version = await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=raw_artifact,
|
|
)
|
|
assert version == 0
|
|
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.inline_data is not None
|
|
assert loaded.inline_data.mime_type == "text/plain"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_load_artifact_preserves_inline_data_display_name(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Binary artifact load restores inline_data.display_name after save."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
app_name = "app0"
|
|
user_id = "user0"
|
|
session_id = "sess0"
|
|
filename = "artifact.bin"
|
|
display_name = "My Report (final).png"
|
|
artifact = types.Part(
|
|
inline_data=types.Blob(
|
|
mime_type="image/png",
|
|
data=b"\x89PNG\r\n\x1a\n",
|
|
display_name=display_name,
|
|
)
|
|
)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
artifact=artifact,
|
|
)
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name=app_name,
|
|
user_id=user_id,
|
|
session_id=session_id,
|
|
filename=filename,
|
|
)
|
|
|
|
assert loaded is not None
|
|
assert loaded.inline_data is not None
|
|
assert loaded.inline_data.display_name == display_name
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.IN_MEMORY,
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"text_content",
|
|
['{"key": "value"}', "some other text"],
|
|
)
|
|
async def test_save_load_text_artifact(
|
|
service_type, artifact_service_factory, text_content
|
|
):
|
|
"""Tests that text artifacts retain .text after round-trip save/load."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
artifact = types.Part.from_text(text=text_content)
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="data.json",
|
|
artifact=artifact,
|
|
)
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="data.json",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.text == text_content
|
|
assert loaded.inline_data is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"service_type",
|
|
[
|
|
ArtifactServiceType.GCS,
|
|
ArtifactServiceType.FILE,
|
|
],
|
|
)
|
|
async def test_save_load_empty_text_artifact(
|
|
service_type, artifact_service_factory
|
|
):
|
|
"""Tests that empty text artifacts survive round-trip save/load."""
|
|
artifact_service = artifact_service_factory(service_type)
|
|
artifact = types.Part.from_text(text="")
|
|
|
|
await artifact_service.save_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="empty.txt",
|
|
artifact=artifact,
|
|
)
|
|
loaded = await artifact_service.load_artifact(
|
|
app_name="app0",
|
|
user_id="user0",
|
|
session_id="123",
|
|
filename="empty.txt",
|
|
)
|
|
assert loaded is not None
|
|
assert loaded.text == ""
|
|
assert loaded.inline_data is None
|
|
|
|
|
|
def test_file_uri_to_path_normalizes_windows_file_uri(monkeypatch):
|
|
monkeypatch.setattr(file_artifact_service, "os", SimpleNamespace(name="nt"))
|
|
mocked_url2pathname = mock.Mock(return_value=r"C:\tmp\adk artifacts")
|
|
monkeypatch.setattr(
|
|
file_artifact_service, "url2pathname", mocked_url2pathname
|
|
)
|
|
|
|
result = file_artifact_service._file_uri_to_path(
|
|
"file:///C:/tmp/adk%20artifacts"
|
|
)
|
|
|
|
mocked_url2pathname.assert_called_once_with("/C:/tmp/adk artifacts")
|
|
assert result == Path(r"C:\tmp\adk artifacts")
|
|
|
|
|
|
def test_file_uri_to_path_returns_none_for_non_file_uri():
|
|
assert (
|
|
file_artifact_service._file_uri_to_path("gs://bucket/adk_artifacts")
|
|
is None
|
|
)
|