tests_socket.py

import pytest
import asyncio
from fastapi.testclient import TestClient
from fastapi import FastAPI, WebSocket
from uuid import UUID, uuid4
from datetime import datetime, timedelta, timezone
import json
from unittest.mock import Mock, patch

from channel.socket_api import router as socket_router
from channel.socket_service import SocketService, CONNECTION_TIMEOUT, HEARTBEAT_INTERVAL

# Test app setup
app = FastAPI()
app.include_router(socket_router)

@pytest.fixture
def socket_service():
    return SocketService()

@pytest.fixture
def test_client():
    return TestClient(app)

@pytest.fixture
def user_id():
    return uuid4()

@pytest.fixture
def device_id():
    return "test_device_1"

class MockWebSocket:
    def __init__(self):
        self.sent_messages = []
        self.closed = False
        
    async def accept(self):
        pass
        
    async def send_json(self, message):
        self.sent_messages.append(message)
        
    async def receive_json(self):
        return {"type": "heartbeat_response"}
        
    async def close(self):
        self.closed = True

@pytest.mark.asyncio
async def test_websocket_connection(socket_service):
    # Test basic connection
    ws = MockWebSocket()
    user_id = uuid4()
    device_id = "test_device"
    
    await socket_service.connect(ws, user_id, device_id)
    
    connection_key = f"socket:{user_id}:{device_id}"
    assert connection_key in socket_service.active_connections
    assert connection_key in socket_service.connection_times

@pytest.mark.asyncio
async def test_multiple_device_connections(socket_service):
    # Test multiple devices for same user
    user_id = uuid4()
    devices = ["device1", "device2", "device3"]
    
    for device_id in devices:
        ws = MockWebSocket()
        await socket_service.connect(ws, user_id, device_id)
    
    user_connections = await socket_service.get_user_connections(user_id)
    assert len(user_connections) == len(devices)
    
    for device_id in devices:
        assert device_id in user_connections

@pytest.mark.asyncio
async def test_broadcast_message(socket_service):
    # Test broadcasting messages to multiple devices
    user_id = uuid4()
    devices = ["device1", "device2", "device3"]
    websockets = {}
    
    # Connect multiple devices
    for device_id in devices:
        ws = MockWebSocket()
        websockets[device_id] = ws
        await socket_service.connect(ws, user_id, device_id)
    
    # Test broadcast
    test_message = {"type": "test", "content": "Hello, devices!"}
    broadcast_count = await socket_service.broadcast_message(user_id, test_message)
    
    assert broadcast_count == len(devices)
    
    # Verify each device received the message
    for ws in websockets.values():
        assert len(ws.sent_messages) == 1
        assert ws.sent_messages[0] == test_message

@pytest.mark.asyncio
async def test_broadcast_with_exclusion(socket_service):
    # Test broadcasting with device exclusion
    user_id = uuid4()
    devices = ["device1", "device2", "device3"]
    websockets = {}
    
    for device_id in devices:
        ws = MockWebSocket()
        websockets[device_id] = ws
        await socket_service.connect(ws, user_id, device_id)
    
    test_message = {"type": "test", "content": "Hello, except device1!"}
    exclude_device = "device1"
    broadcast_count = await socket_service.broadcast_message(user_id, test_message, exclude_device)
    
    assert broadcast_count == len(devices) - 1
    assert len(websockets[exclude_device].sent_messages) == 0
    
    for device_id, ws in websockets.items():
        if device_id != exclude_device:
            assert len(ws.sent_messages) == 1
            assert ws.sent_messages[0] == test_message

@pytest.mark.asyncio
async def test_connection_cleanup(socket_service):
    # Test automatic cleanup of expired connections
    user_id = uuid4()
    device_id = "test_device"
    ws = MockWebSocket()
    
    await socket_service.connect(ws, user_id, device_id)
    connection_key = f"socket:{user_id}:{device_id}"
    
    # Simulate connection expiration
    socket_service.connection_times[connection_key] = datetime.now(timezone.utc) - timedelta(seconds=CONNECTION_TIMEOUT + 1)
    
    # Run cleanup for 3 iterations with a short interval
    await socket_service.cleanup_expired_connections(iterations=3, cleanup_interval=0.1)
    
    assert connection_key not in socket_service.active_connections
    assert connection_key not in socket_service.connection_times
    assert ws.closed is True

@pytest.mark.asyncio
async def test_cleanup_multiple_iterations(socket_service):
    # Test cleanup over multiple iterations
    user_id = uuid4()
    devices = ["device1", "device2", "device3"]
    
    # Connect three devices
    for device_id in devices:
        ws = MockWebSocket()
        await socket_service.connect(ws, user_id, device_id)
        
    # Expire connections at different times
    for i, device_id in enumerate(devices):
        connection_key = f"socket:{user_id}:{device_id}"
        if i < 2:  # First two devices are expired
            # First device expired long ago, second just expired
            expiry_time = datetime.now(timezone.utc) - timedelta(seconds=CONNECTION_TIMEOUT + (i * 100))
            socket_service.connection_times[connection_key] = expiry_time
        # Third device keeps its current timestamp (active)
    
    # Run cleanup for 3 iterations with a short interval
    await socket_service.cleanup_expired_connections(iterations=3, cleanup_interval=0.1)
    
    # Check results
    remaining_connections = await socket_service.get_user_connections(user_id)
    assert len(remaining_connections) == 1  # Only the last device should remain
    assert "device3" in remaining_connections  # The non-expired connection

@pytest.mark.asyncio
async def test_heartbeat_update(socket_service):
    # Test heartbeat mechanism
    user_id = uuid4()
    device_id = "test_device"
    ws = MockWebSocket()
    
    await socket_service.connect(ws, user_id, device_id)
    connection_key = f"socket:{user_id}:{device_id}"
    
    initial_time = socket_service.connection_times[connection_key]
    await asyncio.sleep(0.1)  # Small delay
    
    await socket_service.update_heartbeat(user_id, device_id)
    updated_time = socket_service.connection_times[connection_key]
    
    assert updated_time > initial_time

@pytest.mark.asyncio
async def test_disconnect(socket_service):
    # Test disconnection
    user_id = uuid4()
    device_id = "test_device"
    ws = MockWebSocket()
    
    await socket_service.connect(ws, user_id, device_id)
    connection_key = f"socket:{user_id}:{device_id}"
    
    assert connection_key in socket_service.active_connections
    
    await socket_service.disconnect(user_id, device_id)
    
    assert connection_key not in socket_service.active_connections
    assert connection_key not in socket_service.connection_times
    assert ws.closed is True