Spaces:
Running
Running
File size: 2,825 Bytes
0772b5a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | 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
@pytest.mark.asyncio
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()
|