image_management_api/tests/services/test_storage_service.py

import os
import pytest
import uuid
from fastapi import UploadFile
from unittest.mock import patch, MagicMock
from io import BytesIO

from src.services.storage import StorageService
from src.db.repositories.image_repository import ImageRepository, image_repository
from src.db.models.image import ImageModel

# Hardcoded API key as requested
API_KEY = "Wwg4eJjJ.d03970d43cf3a454ad4168b3226b423f"

# Mock team ID for testing
MOCK_TEAM_ID = "test-team-123"

@pytest.fixture
def test_image_path():
    """Get path to test image"""
    # Assuming image.png exists in the images directory
    return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "images", "image.png")

@pytest.fixture
def test_image_data(test_image_path):
    """Get test image data"""
    with open(test_image_path, "rb") as f:
        return f.read()

@pytest.fixture
def test_upload_file(test_image_data):
    """Create a test UploadFile object"""
    file = UploadFile(
        filename="test_image.png",
        file=BytesIO(test_image_data),
        content_type="image/png"
    )
    return file

@pytest.mark.asyncio
async def test_upload_image_and_verify():
    """Test uploading an image and verifying it was added to storage and database"""
    # Create mocks
    mock_storage_client = MagicMock()
    mock_bucket = MagicMock()
    mock_blob = MagicMock()
    
    # Configure mocks
    mock_storage_client.bucket.return_value = mock_bucket
    mock_bucket.exists.return_value = True
    mock_bucket.blob.return_value = mock_blob
    mock_blob.exists.return_value = True
    
    # Generate a unique filename for the test
    test_filename = f"test-{uuid.uuid4().hex}.png"
    test_content = b"test image content"
    test_content_type = "image/png"
    test_file_size = len(test_content)
    
    # Create a test upload file
    upload_file = UploadFile(
        filename=test_filename,
        file=BytesIO(test_content),
        content_type=test_content_type
    )
    
    # Patch the storage client
    with patch('src.services.storage.StorageService._create_storage_client', return_value=mock_storage_client), \
         patch('src.services.storage.StorageService._get_or_create_bucket', return_value=mock_bucket), \
         patch('src.db.repositories.image_repository.ImageRepository.create') as mock_create:
        
        # Configure the mock to return a valid image model
        storage_path = f"{MOCK_TEAM_ID}/{test_filename}"
        mock_image = ImageModel(
            filename=test_filename,
            original_filename=test_filename,
            file_size=test_file_size,
            content_type=test_content_type,
            storage_path=storage_path,
            team_id=MOCK_TEAM_ID,
            uploader_id="test-user-123"
        )
        mock_create.return_value = mock_image
        
        # Create a storage service instance
        storage_service = StorageService()
        
        # Upload the file
        storage_path, content_type, file_size, metadata = await storage_service.upload_file(
            upload_file, MOCK_TEAM_ID
        )
        
        # Verify the file was uploaded
        mock_bucket.blob.assert_called_with(storage_path)
        mock_blob.upload_from_string.assert_called_once()
        
        # Verify storage path and content type
        assert storage_path.startswith(MOCK_TEAM_ID)
        assert content_type == test_content_type
        assert file_size == test_file_size

@pytest.mark.asyncio
async def test_upload_and_retrieve_image():
    """Test uploading an image and then retrieving it"""
    # Create mocks
    mock_storage_client = MagicMock()
    mock_bucket = MagicMock()
    mock_blob = MagicMock()
    
    # Configure mocks
    mock_storage_client.bucket.return_value = mock_bucket
    mock_bucket.exists.return_value = True
    mock_bucket.blob.return_value = mock_blob
    mock_blob.exists.return_value = True
    
    # Set up the blob to return test content when downloaded
    test_content = b"test image content"
    mock_blob.download_as_bytes.return_value = test_content
    
    # Generate a unique filename for the test
    test_filename = f"test-{uuid.uuid4().hex}.png"
    test_content_type = "image/png"
    test_file_size = len(test_content)
    
    # Create a test upload file
    upload_file = UploadFile(
        filename=test_filename,
        file=BytesIO(test_content),
        content_type=test_content_type
    )
    
    # Patch the storage client
    with patch('src.services.storage.StorageService._create_storage_client', return_value=mock_storage_client), \
         patch('src.services.storage.StorageService._get_or_create_bucket', return_value=mock_bucket):
        
        # Create a storage service instance
        storage_service = StorageService()
        
        # Upload the file
        storage_path, content_type, file_size, metadata = await storage_service.upload_file(
            upload_file, MOCK_TEAM_ID
        )
        
        # Retrieve the file
        retrieved_content = storage_service.get_file(storage_path)
        
        # Verify the content was retrieved
        mock_blob.download_as_bytes.assert_called_once()
        assert retrieved_content == test_content

@pytest.mark.asyncio
async def test_upload_with_real_image(test_upload_file):
    """Test uploading a real image file (requires proper configuration)"""
    # This test requires actual credentials and configuration
    # It will be skipped by default to avoid failures in CI environments
    
    # Skip if no API key environment is configured
    if not API_KEY:
        pytest.skip("API key not configured")
    
    # Create mocks for storage
    mock_storage_client = MagicMock()
    mock_bucket = MagicMock()
    mock_blob = MagicMock()
    
    # Configure mocks
    mock_storage_client.bucket.return_value = mock_bucket
    mock_bucket.exists.return_value = True
    mock_bucket.blob.return_value = mock_blob
    
    # Patch the storage client
    with patch('src.services.storage.StorageService._create_storage_client', return_value=mock_storage_client), \
         patch('src.services.storage.StorageService._get_or_create_bucket', return_value=mock_bucket), \
         patch('src.db.repositories.image_repository.ImageRepository.create') as mock_create:
        
        # Create a storage service instance
        storage_service = StorageService()
        
        # Upload the real image file
        storage_path, content_type, file_size, metadata = await storage_service.upload_file(
            test_upload_file, MOCK_TEAM_ID
        )
        
        # Verify the file was uploaded
        assert storage_path.startswith(MOCK_TEAM_ID)
        assert content_type == "image/png"
        assert file_size > 0
        
        # Check if metadata was extracted (if it's an image)
        if content_type.startswith('image/'):
            assert 'width' in metadata
            assert 'height' in metadata