Spaces:

achin42
/

channel

Build error

channel / tests_socket.py

Upload folder using huggingface_hub

bc1e6dc verified almost 2 years ago

6.98 kB

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 

	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