File size: 6,979 Bytes
bc1e6dc | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | 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 |