def get_artifact_service(
    service_type: ArtifactServiceType = ArtifactServiceType.IN_MEMORY,
):
  """Creates an artifact service for testing."""
  if service_type == ArtifactServiceType.GCS:
    return mock_gcs_artifact_service()
  return InMemoryArtifactService()
Creates an artifact service for testing.
tests/unittests/artifacts/test_artifact_service.py