wehub-resource-sync
198 lines
7.1 KiB
Python
198 lines
7.1 KiB
Python
# Copyright 2025 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.
|
|
|
|
"""Tests for the search controller and service."""
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
from unittest.mock import MagicMock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from google.genai import types
|
|
from PIL import Image as PIL_Image
|
|
|
|
from src.controller.search import router
|
|
from src.model.search import (
|
|
CreateSearchRequest,
|
|
ImageGenerationResult,
|
|
CustomImageResult,
|
|
)
|
|
from src.service.search import ImagenSearchService
|
|
|
|
client = TestClient(router)
|
|
|
|
|
|
def create_minimal_image_bytes(image_format="PNG") -> bytes:
|
|
"""Creates bytes for a 1x1 pixel image."""
|
|
img = PIL_Image.new("RGB", (1, 1), color="red")
|
|
buffer = BytesIO()
|
|
img.save(buffer, format=image_format)
|
|
return buffer.getvalue()
|
|
|
|
|
|
@pytest.fixture(scope="function", name="mock_genai_client")
|
|
def fixture_mock_genai_client():
|
|
"""Provides a mock google.genai Client."""
|
|
mock_client = MagicMock()
|
|
mock_response = types.EditImageResponse()
|
|
|
|
# Create mock generated images with base64 encoded placeholder image data
|
|
mock_image_data = types.Image(
|
|
gcs_uri="gs://mock_bucket/mock_image.png",
|
|
image_bytes=b"mock_image_bytes", # Must be bytes
|
|
mime_type="image/png",
|
|
)
|
|
mock_generated_image = types.GeneratedImage(
|
|
enhanced_prompt="Mock enhanced prompt",
|
|
rai_filtered_reason=None,
|
|
image=mock_image_data,
|
|
)
|
|
mock_response.generated_images = [
|
|
mock_generated_image,
|
|
mock_generated_image,
|
|
mock_generated_image,
|
|
mock_generated_image,
|
|
]
|
|
|
|
mock_client.models.edit_image.return_value = mock_response
|
|
return mock_client
|
|
|
|
|
|
class TestSearchController:
|
|
"""Tests for the /api/search endpoint."""
|
|
|
|
def test_search_endpoint(self, monkeypatch, mock_genai_client):
|
|
# Mock the ImagenSearchService to avoid actual API calls
|
|
# Mock the google.auth.default to avoid authentication issues
|
|
with monkeypatch.context() as m:
|
|
mock_client_class = MagicMock(return_value=mock_genai_client)
|
|
m.setattr("src.service.search.genai.Client", mock_client_class)
|
|
m.setattr(
|
|
"src.service.search.google.auth.default",
|
|
lambda: (None, "test_project_id"),
|
|
)
|
|
# Mock OpenCV functions used within the service
|
|
mock_cascade = MagicMock()
|
|
mock_cascade.detectMultiScale.return_value = [
|
|
(10, 10, 50, 50)
|
|
] # Mock finding one face
|
|
m.setattr(
|
|
"src.service.search.cv2.CascadeClassifier",
|
|
lambda x: mock_cascade,
|
|
)
|
|
# Mock imdecode to return a dummy numpy array with a shape
|
|
mock_img_array = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
m.setattr(
|
|
"src.service.search.cv2.imdecode",
|
|
lambda buf, flags: mock_img_array,
|
|
)
|
|
# Mock cvtColor to return a dummy grayscale array
|
|
mock_gray_array = np.zeros((100, 100), dtype=np.uint8)
|
|
m.setattr(
|
|
"src.service.search.cv2.cvtColor",
|
|
lambda img, code: mock_gray_array,
|
|
)
|
|
|
|
search_term = "a cute cat wearing a hat"
|
|
image_content = create_minimal_image_bytes()
|
|
user_image = ("test_image.png", BytesIO(image_content), "image/png")
|
|
|
|
response = client.post(
|
|
"/api/search",
|
|
data={
|
|
"term": search_term,
|
|
"numberOfImages": 4,
|
|
"maskDistilation": 0.005,
|
|
"generationModel": "imagen-3.0-capability-001",
|
|
},
|
|
files={"userImage": user_image},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data) == 8
|
|
|
|
for image_data in data:
|
|
assert image_data["enhancedPrompt"] == "Mock enhanced prompt"
|
|
assert (
|
|
image_data["image"]["gcsUri"]
|
|
== "gs://mock_bucket/mock_image.png"
|
|
)
|
|
assert image_data["image"]["mimeType"] == "image/png"
|
|
assert image_data["image"]["encodedImage"] == base64.b64encode(
|
|
b"mock_image_bytes"
|
|
).decode("utf-8")
|
|
|
|
|
|
class TestImagenSearchService:
|
|
"""Tests for the ImagenSearchService class."""
|
|
|
|
def test_imagen_search_service(self, monkeypatch, mock_genai_client):
|
|
|
|
# Mock the google.auth.default to avoid authentication issues
|
|
with monkeypatch.context() as m:
|
|
mock_client_class = MagicMock(return_value=mock_genai_client)
|
|
m.setattr("src.service.search.genai.Client", mock_client_class)
|
|
m.setattr(
|
|
"src.service.search.google.auth.default",
|
|
lambda: (None, "test_project_id"),
|
|
)
|
|
# Mock OpenCV functions used within the service
|
|
mock_cascade = MagicMock()
|
|
mock_cascade.detectMultiScale.return_value = [
|
|
(10, 10, 50, 50)
|
|
] # Mock finding one face
|
|
m.setattr(
|
|
"src.service.search.cv2.CascadeClassifier",
|
|
lambda x: mock_cascade,
|
|
)
|
|
# Mock imdecode to return a dummy numpy array with a shape
|
|
mock_img_array = np.zeros((100, 100, 3), dtype=np.uint8)
|
|
m.setattr(
|
|
"src.service.search.cv2.imdecode",
|
|
lambda buf, flags: mock_img_array,
|
|
)
|
|
# Mock cvtColor to return a dummy grayscale array
|
|
mock_gray_array = np.zeros((100, 100), dtype=np.uint8)
|
|
m.setattr(
|
|
"src.service.search.cv2.cvtColor",
|
|
lambda img, code: mock_gray_array,
|
|
)
|
|
|
|
valid_image_bytes = create_minimal_image_bytes()
|
|
search_request = CreateSearchRequest(
|
|
term="a dog playing fetch",
|
|
user_image=valid_image_bytes,
|
|
number_of_images=2,
|
|
mask_distilation=0.1,
|
|
generation_model="imagen-3.0-capability-001",
|
|
)
|
|
service = ImagenSearchService()
|
|
results = service.generate_images(search_request)
|
|
|
|
assert isinstance(results, list)
|
|
assert len(results) == 8
|
|
assert all(
|
|
isinstance(result, ImageGenerationResult) for result in results
|
|
)
|
|
mock_client_class.assert_called_once()
|
|
|
|
for result in results:
|
|
assert isinstance(result.image, CustomImageResult)
|
|
assert result.image.encoded_image == base64.b64encode(
|
|
b"mock_image_bytes"
|
|
).decode("utf-8")
|