test_artifact_service.py

# 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 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.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
@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_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,
    )


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 loads an internal artifact:// URI reference."""
  service = mock_gcs_artifact_service()  # type: ignore[no-untyped-call]
  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.file_data is not None
  assert loaded.file_data.file_uri == artifact_ref_uri


@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
@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).

  Regression test for https://github.com/google/adk-python/issues/2886
  """
  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"