Spaces:
Running
Running
| import pytest | |
| from unittest.mock import MagicMock, patch | |
| import uuid | |
| from fastapi import HTTPException | |
| from app.chat_image_upload import cleanup_session_storage, validate_chat_image_bytes | |
| from app.session_cache import ( | |
| invalidate_session_owner, | |
| session_owned_by_user, | |
| ) | |
| from app.websocket_manager import active_connections, notify_status | |
| def test_session_owner_authoritative_check(): | |
| user_id = str(uuid.uuid4()) | |
| session_id = str(uuid.uuid4()) | |
| fetch_count = 0 | |
| def mock_owns(): | |
| nonlocal fetch_count | |
| fetch_count += 1 | |
| return True | |
| assert session_owned_by_user(session_id, user_id, mock_owns) is True | |
| assert fetch_count == 1 | |
| invalidate_session_owner(session_id, user_id) | |
| assert session_owned_by_user(session_id, user_id, mock_owns) is True | |
| assert fetch_count == 2 | |
| def test_chat_image_validation_magic_bytes(): | |
| # Valid PNG magic bytes | |
| png_bytes = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01" | |
| ext, mime = validate_chat_image_bytes("test.png", png_bytes, "image/png") | |
| assert ext == ".png" | |
| assert mime == "image/png" | |
| # Valid JPEG magic bytes | |
| jpeg_bytes = b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01" | |
| ext, mime = validate_chat_image_bytes("test.jpg", jpeg_bytes, "image/jpeg") | |
| assert ext == ".jpg" | |
| assert mime == "image/jpeg" | |
| # Invalid corrupted content | |
| bad_bytes = b"NOT_AN_IMAGE_DATA_HEADER" | |
| with pytest.raises(HTTPException) as exc_info: | |
| validate_chat_image_bytes("test.png", bad_bytes, "image/png") | |
| assert exc_info.value.status_code == 400 | |
| def test_cleanup_session_storage(): | |
| session_id = str(uuid.uuid4()) | |
| mock_supabase = MagicMock() | |
| mock_from = MagicMock() | |
| mock_from.list.return_value = [ | |
| {"name": f"image_v1_{session_id}.png"}, | |
| {"name": ".emptyFolderPlaceholder"}, | |
| ] | |
| mock_supabase.storage.from_.return_value = mock_from | |
| with patch("app.supabase_client.get_supabase", return_value=mock_supabase): | |
| cleanup_session_storage(session_id) | |
| assert mock_supabase.storage.from_.called | |
| assert mock_from.remove.called | |
| async def test_websocket_dead_connection_pruning(): | |
| job_id = str(uuid.uuid4()) | |
| from unittest.mock import AsyncMock | |
| good_ws = MagicMock() | |
| good_ws.send_json = AsyncMock() | |
| bad_ws = MagicMock() | |
| async def bad_send(_): | |
| raise RuntimeError("Connection closed") | |
| bad_ws.send_json = bad_send | |
| active_connections[job_id] = [good_ws, bad_ws] | |
| # Notify status should prune bad_ws | |
| await notify_status(job_id, {"status": "processing"}) | |
| assert job_id in active_connections | |
| assert bad_ws not in active_connections[job_id] | |
| assert good_ws in active_connections[job_id] | |
| # Clean up | |
| active_connections.clear() | |