424 lines
16 KiB
Python
424 lines
16 KiB
Python
import pytest
|
|
from unittest.mock import Mock, patch, AsyncMock
|
|
from fastapi.testclient import TestClient
|
|
from fastapi import status
|
|
import io
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
from bson import ObjectId as PyObjectId
|
|
|
|
from src.api.v1.images import router
|
|
from src.models.user import UserModel
|
|
from src.models.team import PyObjectId
|
|
|
|
|
|
class TestImageUploadWithPubSub:
|
|
"""Test cases for image upload with Pub/Sub integration"""
|
|
|
|
@pytest.fixture
|
|
def mock_current_user(self):
|
|
"""Mock current user for authentication"""
|
|
user = Mock()
|
|
user.id = PyObjectId()
|
|
user.team_id = PyObjectId()
|
|
user.email = "test@example.com"
|
|
user.name = "Test User"
|
|
user.is_admin = True
|
|
user.is_active = True
|
|
return user
|
|
|
|
@pytest.fixture
|
|
def test_image_file(self):
|
|
"""Create a test image file"""
|
|
# Create a simple test image file
|
|
image_data = b"fake image data for testing"
|
|
return BytesIO(image_data)
|
|
|
|
@pytest.fixture
|
|
def mock_storage_service(self):
|
|
"""Mock storage service"""
|
|
with patch('src.api.v1.images.storage_service') as mock_service:
|
|
mock_service.upload_file = AsyncMock(return_value=(
|
|
"bucket/team/test-image.jpg", # storage_path
|
|
"image/jpeg", # content_type
|
|
1024, # file_size
|
|
{"width": 800, "height": 600} # metadata
|
|
))
|
|
yield mock_service
|
|
|
|
@pytest.fixture
|
|
def mock_image_repository(self):
|
|
"""Mock image repository"""
|
|
with patch('src.api.v1.images.image_repository') as mock_repo:
|
|
# Create a mock image object
|
|
mock_image = Mock()
|
|
mock_image.id = PyObjectId()
|
|
mock_image.filename = "test-image-123.jpg"
|
|
mock_image.original_filename = "test.jpg"
|
|
mock_image.file_size = 1024
|
|
mock_image.content_type = "image/jpeg"
|
|
mock_image.storage_path = "bucket/team/test-image.jpg"
|
|
mock_image.team_id = PyObjectId()
|
|
mock_image.uploader_id = PyObjectId()
|
|
mock_image.description = None
|
|
mock_image.metadata = {}
|
|
mock_image.upload_date = Mock()
|
|
mock_image.has_embedding = False
|
|
|
|
# Configure the create method to return the mock image
|
|
mock_repo.create = AsyncMock(return_value=mock_image)
|
|
yield mock_repo
|
|
|
|
@pytest.fixture
|
|
def mock_pubsub_service(self):
|
|
"""Mock Pub/Sub service"""
|
|
with patch('src.api.v1.images.pubsub_service') as mock_service:
|
|
mock_service.publish_image_processing_task = AsyncMock(return_value=True)
|
|
yield mock_service
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_image_publishes_to_pubsub(
|
|
self,
|
|
mock_current_user,
|
|
test_image_file,
|
|
mock_storage_service,
|
|
mock_image_repository,
|
|
mock_pubsub_service
|
|
):
|
|
"""Test that image upload publishes a message to Pub/Sub"""
|
|
with patch('src.api.v1.images.get_current_user', return_value=mock_current_user):
|
|
from src.api.v1.images import upload_image
|
|
from fastapi import UploadFile
|
|
|
|
# Create upload file
|
|
upload_file = UploadFile(
|
|
filename="test.jpg",
|
|
file=test_image_file,
|
|
content_type="image/jpeg"
|
|
)
|
|
|
|
# Mock request
|
|
request = Mock()
|
|
request.url.path = "/api/v1/images"
|
|
request.method = "POST"
|
|
|
|
# Call the upload function
|
|
response = await upload_image(
|
|
request=request,
|
|
file=upload_file,
|
|
current_user=mock_current_user
|
|
)
|
|
|
|
# Verify storage service was called
|
|
mock_storage_service.upload_file.assert_called_once()
|
|
|
|
# Verify image was created in repository
|
|
mock_image_repository.create.assert_called_once()
|
|
|
|
# Verify Pub/Sub task was published
|
|
mock_pubsub_service.publish_image_processing_task.assert_called_once()
|
|
|
|
# Get the call arguments for Pub/Sub
|
|
pubsub_call_args = mock_pubsub_service.publish_image_processing_task.call_args
|
|
task_data = pubsub_call_args[0][0] # First positional argument
|
|
|
|
# Verify task data contains expected fields
|
|
assert "image_id" in task_data
|
|
assert "team_id" in task_data
|
|
assert "storage_path" in task_data
|
|
assert "content_type" in task_data
|
|
|
|
# Verify response
|
|
assert response.filename == "test-image-123.jpg"
|
|
assert response.content_type == "image/jpeg"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_image_pubsub_failure_continues(
|
|
self,
|
|
mock_current_user,
|
|
test_image_file,
|
|
mock_storage_service,
|
|
mock_image_repository,
|
|
mock_pubsub_service
|
|
):
|
|
"""Test that image upload continues even if Pub/Sub fails"""
|
|
with patch('src.api.v1.images.get_current_user', return_value=mock_current_user):
|
|
from src.api.v1.images import upload_image
|
|
from fastapi import UploadFile
|
|
|
|
# Configure Pub/Sub to fail
|
|
mock_pubsub_service.publish_image_processing_task = AsyncMock(return_value=False)
|
|
|
|
# Create upload file
|
|
upload_file = UploadFile(
|
|
filename="test.jpg",
|
|
file=test_image_file,
|
|
content_type="image/jpeg"
|
|
)
|
|
|
|
# Mock request
|
|
request = Mock()
|
|
request.url.path = "/api/v1/images"
|
|
request.method = "POST"
|
|
|
|
# Call the upload function
|
|
response = await upload_image(
|
|
request=request,
|
|
file=upload_file,
|
|
current_user=mock_current_user
|
|
)
|
|
|
|
# Verify storage and repository were still called
|
|
mock_storage_service.upload_file.assert_called_once()
|
|
mock_image_repository.create.assert_called_once()
|
|
|
|
# Verify Pub/Sub was attempted
|
|
mock_pubsub_service.publish_image_processing_task.assert_called_once()
|
|
|
|
# Verify response is still successful
|
|
assert response.filename == "test-image-123.jpg"
|
|
assert response.content_type == "image/jpeg"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_image_pubsub_exception_continues(
|
|
self,
|
|
mock_current_user,
|
|
test_image_file,
|
|
mock_storage_service,
|
|
mock_image_repository,
|
|
mock_pubsub_service
|
|
):
|
|
"""Test that image upload continues even if Pub/Sub raises an exception"""
|
|
with patch('src.api.v1.images.get_current_user', return_value=mock_current_user):
|
|
from src.api.v1.images import upload_image
|
|
from fastapi import UploadFile
|
|
|
|
# Configure Pub/Sub to raise an exception
|
|
mock_pubsub_service.publish_image_processing_task = AsyncMock(
|
|
side_effect=Exception("Pub/Sub service unavailable")
|
|
)
|
|
|
|
# Create upload file
|
|
upload_file = UploadFile(
|
|
filename="test.jpg",
|
|
file=test_image_file,
|
|
content_type="image/jpeg"
|
|
)
|
|
|
|
# Mock request
|
|
request = Mock()
|
|
request.url.path = "/api/v1/images"
|
|
request.method = "POST"
|
|
|
|
# Call the upload function
|
|
response = await upload_image(
|
|
request=request,
|
|
file=upload_file,
|
|
current_user=mock_current_user
|
|
)
|
|
|
|
# Verify storage and repository were still called
|
|
mock_storage_service.upload_file.assert_called_once()
|
|
mock_image_repository.create.assert_called_once()
|
|
|
|
# Verify Pub/Sub was attempted
|
|
mock_pubsub_service.publish_image_processing_task.assert_called_once()
|
|
|
|
# Verify response is still successful
|
|
assert response.filename == "test.jpg"
|
|
assert response.content_type == "image/jpeg"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_image_with_description(
|
|
self,
|
|
mock_current_user,
|
|
test_image_file,
|
|
mock_storage_service,
|
|
mock_image_repository,
|
|
mock_pubsub_service
|
|
):
|
|
"""Test image upload with description"""
|
|
with patch('src.api.v1.images.get_current_user', return_value=mock_current_user):
|
|
from src.api.v1.images import upload_image
|
|
from fastapi import UploadFile
|
|
|
|
# Create upload file
|
|
upload_file = UploadFile(
|
|
filename="test.jpg",
|
|
file=test_image_file,
|
|
content_type="image/jpeg"
|
|
)
|
|
|
|
# Mock request
|
|
request = Mock()
|
|
request.url.path = "/api/v1/images"
|
|
request.method = "POST"
|
|
|
|
# Call the upload function with description
|
|
response = await upload_image(
|
|
request=request,
|
|
file=upload_file,
|
|
description="Test image",
|
|
current_user=mock_current_user
|
|
)
|
|
|
|
# Verify Pub/Sub task was published
|
|
mock_pubsub_service.publish_image_processing_task.assert_called_once()
|
|
|
|
# Verify image was created with correct data
|
|
mock_image_repository.create.assert_called_once()
|
|
created_image_data = mock_image_repository.create.call_args[0][0]
|
|
|
|
assert created_image_data.description == "Test image"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_image_with_collection_id(
|
|
self,
|
|
mock_current_user,
|
|
test_image_file,
|
|
mock_storage_service,
|
|
mock_image_repository,
|
|
mock_pubsub_service
|
|
):
|
|
"""Test image upload with collection ID"""
|
|
with patch('src.api.v1.images.get_current_user', return_value=mock_current_user):
|
|
from src.api.v1.images import upload_image
|
|
from fastapi import UploadFile
|
|
|
|
# Create upload file
|
|
upload_file = UploadFile(
|
|
filename="test.jpg",
|
|
file=test_image_file,
|
|
content_type="image/jpeg"
|
|
)
|
|
|
|
# Mock request
|
|
request = Mock()
|
|
request.url.path = "/api/v1/images"
|
|
request.method = "POST"
|
|
|
|
collection_id = str(PyObjectId())
|
|
|
|
# Call the upload function with collection ID
|
|
response = await upload_image(
|
|
request=request,
|
|
file=upload_file,
|
|
collection_id=collection_id,
|
|
current_user=mock_current_user
|
|
)
|
|
|
|
# Verify Pub/Sub task was published
|
|
mock_pubsub_service.publish_image_processing_task.assert_called_once()
|
|
|
|
# Verify image was created with collection ID
|
|
mock_image_repository.create.assert_called_once()
|
|
created_image_data = mock_image_repository.create.call_args[0][0]
|
|
|
|
assert str(created_image_data.collection_id) == collection_id
|
|
|
|
|
|
class TestImageModelUpdates:
|
|
"""Test cases for updated image model with embedding fields"""
|
|
|
|
def test_image_model_has_embedding_fields(self):
|
|
"""Test that ImageModel has the new embedding fields"""
|
|
from src.models.image import ImageModel
|
|
|
|
# Create an image model instance
|
|
image = ImageModel(
|
|
filename="test.jpg",
|
|
original_filename="test.jpg",
|
|
file_size=1024,
|
|
content_type="image/jpeg",
|
|
storage_path="bucket/path/image.jpg",
|
|
team_id=PyObjectId(),
|
|
uploader_id=PyObjectId()
|
|
)
|
|
|
|
# Check that embedding fields exist with default values
|
|
assert hasattr(image, 'embedding_status')
|
|
assert hasattr(image, 'embedding_error')
|
|
assert hasattr(image, 'embedding_retry_count')
|
|
assert hasattr(image, 'embedding_last_attempt')
|
|
|
|
# Check default values
|
|
assert image.embedding_status == "pending"
|
|
assert image.embedding_error is None
|
|
assert image.embedding_retry_count == 0
|
|
assert image.embedding_last_attempt is None
|
|
assert image.has_embedding is False
|
|
|
|
def test_image_model_embedding_fields_can_be_set(self):
|
|
"""Test that embedding fields can be set"""
|
|
from src.models.image import ImageModel
|
|
from datetime import datetime
|
|
|
|
now = datetime.utcnow()
|
|
|
|
# Create an image model instance with embedding fields
|
|
image = ImageModel(
|
|
filename="test.jpg",
|
|
original_filename="test.jpg",
|
|
file_size=1024,
|
|
content_type="image/jpeg",
|
|
storage_path="bucket/path/image.jpg",
|
|
team_id=PyObjectId(),
|
|
uploader_id=PyObjectId(),
|
|
embedding_status="processing",
|
|
embedding_error="Test error",
|
|
embedding_retry_count=2,
|
|
embedding_last_attempt=now,
|
|
has_embedding=True
|
|
)
|
|
|
|
# Check that values were set correctly
|
|
assert image.embedding_status == "processing"
|
|
assert image.embedding_error == "Test error"
|
|
assert image.embedding_retry_count == 2
|
|
assert image.embedding_last_attempt == now
|
|
assert image.has_embedding is True
|
|
|
|
|
|
class TestPubSubServiceIntegration:
|
|
"""Integration tests for Pub/Sub service with image API"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_end_to_end_image_upload_flow(self):
|
|
"""Test the complete flow from image upload to Pub/Sub message"""
|
|
# This would be an integration test that verifies the entire flow
|
|
# from API call to Pub/Sub message publication
|
|
|
|
# Mock all dependencies
|
|
with patch('src.api.v1.images.storage_service') as mock_storage, \
|
|
patch('src.api.v1.images.image_repository') as mock_repo, \
|
|
patch('src.api.v1.images.pubsub_service') as mock_pubsub, \
|
|
patch('src.api.v1.images.get_current_user') as mock_auth:
|
|
|
|
# Setup mocks
|
|
mock_user = Mock()
|
|
mock_user.id = PyObjectId()
|
|
mock_user.team_id = PyObjectId()
|
|
mock_auth.return_value = mock_user
|
|
|
|
mock_storage.upload_file = AsyncMock(return_value=(
|
|
"bucket/team/image.jpg", "image/jpeg", 1024, {}
|
|
))
|
|
|
|
mock_image = Mock()
|
|
mock_image.id = PyObjectId()
|
|
mock_repo.create = AsyncMock(return_value=mock_image)
|
|
|
|
mock_pubsub.publish_image_processing_task = AsyncMock(return_value=True)
|
|
|
|
# Create test client
|
|
from fastapi import FastAPI
|
|
app = FastAPI()
|
|
app.include_router(router)
|
|
|
|
# This would test the actual HTTP endpoint
|
|
# but requires more complex setup for file uploads
|
|
|
|
# For now, verify the mocks would be called correctly
|
|
assert mock_storage.upload_file is not None
|
|
assert mock_repo.create is not None
|
|
assert mock_pubsub.publish_image_processing_task is not None |