"""Integration tests for download progress WebSocket real-time updates. This module tests the end-to-end flow of download progress from the download service through the WebSocket service to connected clients. """ import asyncio from typing import Any, Dict, List from unittest.mock import Mock, patch import pytest from src.server.models.download import EpisodeIdentifier from src.server.services.anime_service import AnimeService from src.server.services.download_service import DownloadService from src.server.services.progress_service import ProgressService from src.server.services.websocket_service import WebSocketService @pytest.fixture def mock_series_app(): """Mock SeriesApp for testing.""" app = Mock() app.series_list = [] app.search = Mock(return_value=[]) app.ReScan = Mock() def mock_download( serie_folder, season, episode, key, callback=None, **kwargs ): """Simulate download with realistic progress updates.""" if callback: # Simulate yt-dlp progress updates for percent in [10, 25, 50, 75, 90, 100]: callback({ 'percent': float(percent), 'downloaded_mb': percent, 'total_mb': 100.0, 'speed_mbps': 2.5, 'eta_seconds': int((100 - percent) / 2.5), }) result = Mock() result.success = True result.message = "Download completed" return result app.download = Mock(side_effect=mock_download) return app @pytest.fixture def progress_service(): """Create a ProgressService instance.""" return ProgressService() @pytest.fixture def websocket_service(): """Create a WebSocketService instance.""" return WebSocketService() @pytest.fixture async def anime_service(mock_series_app, progress_service): """Create an AnimeService.""" service = AnimeService( series_app=mock_series_app, progress_service=progress_service, ) yield service @pytest.fixture async def download_service(anime_service, progress_service): """Create a DownloadService.""" service = DownloadService( anime_service=anime_service, progress_service=progress_service, persistence_path="/tmp/test_integration_progress_queue.json", ) yield service await service.stop() class TestDownloadProgressIntegration: """Integration tests for download progress WebSocket flow.""" @pytest.mark.asyncio async def test_full_progress_flow_with_websocket( self, download_service, websocket_service ): """Test complete flow from download to WebSocket broadcast.""" # Track all messages sent via WebSocket sent_messages: List[Dict[str, Any]] = [] # Mock WebSocket broadcast methods original_broadcast_progress = ( websocket_service.broadcast_download_progress ) async def mock_broadcast_progress(download_id: str, data: dict): """Capture broadcast calls.""" sent_messages.append({ 'type': 'download_progress', 'download_id': download_id, 'data': data, }) # Call original to maintain functionality await original_broadcast_progress(download_id, data) websocket_service.broadcast_download_progress = ( mock_broadcast_progress ) # Connect download service to WebSocket service async def broadcast_callback(update_type: str, data: dict): """Bridge download service to WebSocket service.""" if update_type == "download_progress": await websocket_service.broadcast_download_progress( data.get("download_id", ""), data, ) download_service.set_broadcast_callback(broadcast_callback) # Add download to queue await download_service.add_to_queue( serie_id="integration_test", serie_folder="test_folder", serie_name="Integration Test Anime", episodes=[EpisodeIdentifier(season=1, episode=1)], ) # Start processing await download_service.start_queue_processing() # Wait for download to complete await asyncio.sleep(1.0) # Verify progress messages were sent progress_messages = [ m for m in sent_messages if m['type'] == 'download_progress' ] assert len(progress_messages) >= 3 # Multiple progress updates # Verify progress increases percentages = [ m['data'].get('progress', {}).get('percent', 0) for m in progress_messages ] # Should have increasing percentages for i in range(1, len(percentages)): assert percentages[i] >= percentages[i - 1] # Last update should be close to 100% assert percentages[-1] >= 90 @pytest.mark.asyncio async def test_websocket_client_receives_progress( self, download_service, websocket_service ): """Test that WebSocket clients receive progress messages.""" # Track messages received by clients client_messages: List[Dict[str, Any]] = [] # Mock WebSocket client class MockWebSocket: """Mock WebSocket for testing.""" async def accept(self): pass async def send_json(self, data): """Capture sent messages.""" client_messages.append(data) async def receive_json(self): # Keep connection open await asyncio.sleep(10) mock_ws = MockWebSocket() # Connect mock client connection_id = "test_client_1" await websocket_service.connect(mock_ws, connection_id) # Connect download service to WebSocket service async def broadcast_callback(update_type: str, data: dict): if update_type == "download_progress": await websocket_service.broadcast_download_progress( data.get("download_id", ""), data, ) download_service.set_broadcast_callback(broadcast_callback) # Add and start download await download_service.add_to_queue( serie_id="client_test", serie_folder="test_folder", serie_name="Client Test Anime", episodes=[EpisodeIdentifier(season=1, episode=1)], ) await download_service.start_queue_processing() await asyncio.sleep(1.0) # Verify client received messages progress_messages = [ m for m in client_messages if m.get('type') == 'download_progress' ] assert len(progress_messages) >= 2 # Cleanup await websocket_service.disconnect(connection_id) @pytest.mark.asyncio async def test_multiple_clients_receive_same_progress( self, download_service, websocket_service ): """Test that all connected clients receive progress updates.""" # Track messages for each client client1_messages: List[Dict] = [] client2_messages: List[Dict] = [] class MockWebSocket: """Mock WebSocket for testing.""" def __init__(self, message_list): self.messages = message_list async def accept(self): pass async def send_json(self, data): self.messages.append(data) async def receive_json(self): await asyncio.sleep(10) # Connect two clients client1 = MockWebSocket(client1_messages) client2 = MockWebSocket(client2_messages) await websocket_service.connect(client1, "client1") await websocket_service.connect(client2, "client2") # Connect download service async def broadcast_callback(update_type: str, data: dict): if update_type == "download_progress": await websocket_service.broadcast_download_progress( data.get("download_id", ""), data, ) download_service.set_broadcast_callback(broadcast_callback) # Start download await download_service.add_to_queue( serie_id="multi_client_test", serie_folder="test_folder", serie_name="Multi Client Test", episodes=[EpisodeIdentifier(season=1, episode=1)], ) await download_service.start_queue_processing() await asyncio.sleep(1.0) # Both clients should receive progress client1_progress = [ m for m in client1_messages if m.get('type') == 'download_progress' ] client2_progress = [ m for m in client2_messages if m.get('type') == 'download_progress' ] assert len(client1_progress) >= 2 assert len(client2_progress) >= 2 # Both should have similar number of updates assert abs(len(client1_progress) - len(client2_progress)) <= 2 # Cleanup await websocket_service.disconnect("client1") await websocket_service.disconnect("client2") @pytest.mark.asyncio async def test_progress_data_structure_matches_frontend_expectations( self, download_service, websocket_service ): """Test that progress data structure matches frontend requirements.""" captured_data: List[Dict] = [] async def capture_broadcast(update_type: str, data: dict): if update_type == "download_progress": captured_data.append(data) await websocket_service.broadcast_download_progress( data.get("download_id", ""), data, ) download_service.set_broadcast_callback(capture_broadcast) await download_service.add_to_queue( serie_id="structure_test", serie_folder="test_folder", serie_name="Structure Test", episodes=[EpisodeIdentifier(season=2, episode=3)], ) await download_service.start_queue_processing() await asyncio.sleep(1.0) assert len(captured_data) > 0 # Verify data structure matches frontend expectations for data in captured_data: # Required fields for frontend (queue.js) assert 'download_id' in data or 'item_id' in data assert 'serie_name' in data assert 'season' in data assert 'episode' in data assert 'progress' in data # Progress object structure progress = data['progress'] assert 'percent' in progress assert 'downloaded_mb' in progress assert 'total_mb' in progress # Verify episode info assert data['season'] == 2 assert data['episode'] == 3 assert data['serie_name'] == "Structure Test" @pytest.mark.asyncio async def test_disconnected_client_doesnt_receive_progress( self, download_service, websocket_service ): """Test that disconnected clients don't receive updates.""" client_messages: List[Dict] = [] class MockWebSocket: async def accept(self): pass async def send_json(self, data): client_messages.append(data) async def receive_json(self): await asyncio.sleep(10) mock_ws = MockWebSocket() # Connect and then disconnect connection_id = "temp_client" await websocket_service.connect(mock_ws, connection_id) await websocket_service.disconnect(connection_id) # Connect download service async def broadcast_callback(update_type: str, data: dict): if update_type == "download_progress": await websocket_service.broadcast_download_progress( data.get("download_id", ""), data, ) download_service.set_broadcast_callback(broadcast_callback) # Start download after disconnect await download_service.add_to_queue( serie_id="disconnect_test", serie_folder="test_folder", serie_name="Disconnect Test", episodes=[EpisodeIdentifier(season=1, episode=1)], ) initial_message_count = len(client_messages) await download_service.start_queue_processing() await asyncio.sleep(1.0) # Should not receive progress updates after disconnect progress_messages = [ m for m in client_messages[initial_message_count:] if m.get('type') == 'download_progress' ] assert len(progress_messages) == 0