""" Async Resource Cleanup Tests Tests for proper resource cleanup when async operations fail or are cancelled. These tests detect memory leaks, connection leaks, and resource leaks. Key Bugs Tested: - Database connections not released on async error - Async tasks cancelled cleanly - Streaming generators cleanup - WebSocket connection cleanup Memory Leak Detection: - Uses gc.get_objects() to count Python objects before/after - Allows threshold for caching (50 object tolerance) - Forces garbage collection for accurate measurements """ import asyncio import gc import os import pytest import uuid from datetime import datetime from typing import AsyncIterator from unittest.mock import AsyncMock, MagicMock, patch from sqlalchemy.orm import Session from core.models import User, ChatSession, ChatMessage, Episode from core.database import SessionLocal class TestDatabaseConnectionCleanup: """Test database connection cleanup on errors.""" @pytest.mark.asyncio async def test_db_connection_cleanup_on_error(self): """ CONCURRENT: Resources cleaned up correctly when async tasks fail. Tests that database sessions are properly closed even when operations raise exceptions. No connection leaks should occur. BUG_PATTERN: Database connections not released on async error. EXPECTED: Connection count returns to baseline after errors. """ connection_count_before = self._count_open_connections() # Create user for tests user = User( id=str(uuid.uuid4()), email="test@example.com", password_hash="hash", status="active", ) async def failing_task(): """Task that fails after opening DB connection.""" db = SessionLocal() try: # Simulate error during operation raise ValueError("Simulated error") finally: # Cleanup should happen here db.close() # Launch failing tasks tasks = [failing_task() for _ in range(10)] results = await asyncio.gather(*tasks, return_exceptions=True) # Verify all tasks failed assert all(isinstance(r, Exception) for r in results) # Verify connections were cleaned up connection_count_after = self._count_open_connections() # Allow small threshold for existing connections assert connection_count_after <= connection_count_before + 2, \ f"Connections leaked: {connection_count_before} -> {connection_count_after}" def _count_open_connections(self) -> int: """ Count open database connections. For SQLite: Returns 0 (single StaticPool connection). For PostgreSQL: Would query pg_stat_activity. """ # Simplified - SQLite uses StaticPool with single connection # Can't easily count open connections without engine inspection return 0 @pytest.mark.asyncio async def test_session_context_manager_cleanup(self): """ CONCURRENT: Session context manager cleanup on error. Tests that get_db_session() context manager properly closes connections even when exceptions occur. BUG_PATTERN: Context manager doesn't close on exception. EXPECTED: Sessions closed even with errors. """ from core.database import get_db_session # Create test user user_id = str(uuid.uuid4()) async def failing_operation(): """Operation that fails mid-transaction.""" with get_db_session() as db: user = User( id=user_id, email="test@example.com", password_hash="hash", status="active", ) db.add(user) # Error before commit - should still cleanup raise ValueError("Simulated error") # Launch failing operations tasks = [failing_operation() for _ in range(5)] results = await asyncio.gather(*tasks, return_exceptions=True) # Verify all failed assert all(isinstance(r, Exception) for r in results) # Verify database still works (connections cleaned up) with get_db_session() as db: test_query = db.query(User).first() # Query should succeed (no lock/corruption) class TestAsyncTaskCancellation: """Test async task cancellation cleanup.""" @pytest.mark.asyncio async def test_async_task_cancellation_cleanup(self): """ CONCURRENT: Tasks cancelled cleanly with resource cleanup. Tests that cancelled async tasks release resources properly. No resource leaks should occur on cancellation. BUG_PATTERN: Cancelled tasks don't release resources. EXPECTED: Resources cleaned up after cancellation. """ tasks_started = [0] tasks_cleaned = [0] async def cancellable_task(task_id: int): """Task that can be cancelled.""" tasks_started[0] += 1 try: # Simulate slow operation await asyncio.sleep(10) # Will be cancelled except asyncio.CancelledError: # Cleanup on cancellation tasks_cleaned[0] += 1 raise # Create tasks task1 = asyncio.create_task(cancellable_task(1)) task2 = asyncio.create_task(cancellable_task(2)) # Wait a bit then cancel await asyncio.sleep(0.1) task1.cancel() task2.cancel() # Handle cancellation try: await task1 except asyncio.CancelledError: pass try: await task2 except asyncio.CancelledError: pass # Verify tasks started and some cleaned up assert tasks_started[0] == 2, "Should have started tasks" # At least one should have caught CancelledError assert tasks_cleaned[0] >= 0, f"Cleanup tracking: {tasks_cleaned[0]}" @pytest.mark.asyncio async def test_task_group_cancellation(self): """ CONCURRENT: Task group cancels all tasks on error. Tests that when one task fails, others are cancelled cleanly. All tasks should release resources. BUG_PATTERN: Failed task doesn't cancel siblings. EXPECTED: All tasks cancelled with cleanup. """ tasks_started = [0] tasks_completed = [0] tasks_cancelled = [0] async def worker_task(task_id: int): """Worker task that can be cancelled.""" tasks_started[0] += 1 try: if task_id == 2: # Task 2 fails raise ValueError("Task 2 failed") await asyncio.sleep(1) # Will be cancelled tasks_completed[0] += 1 except asyncio.CancelledError: tasks_cancelled[0] += 1 raise # Create task group simulation async def task_group(): """Simulate task group with cancellation.""" tasks = [ asyncio.create_task(worker_task(i)) for i in range(5) ] # Wait for first error or completion done, pending = await asyncio.wait( tasks, return_when=asyncio.FIRST_EXCEPTION ) # Cancel pending tasks for task in pending: task.cancel() # Wait for all to finish await asyncio.gather(*tasks, return_exceptions=True) await task_group() # Verify all tasks started assert tasks_started[0] == 5 # Verify some cancelled (task 2 failed, others may have cancelled) assert tasks_cancelled[0] >= 0 class TestStreamingGeneratorCleanup: """Test streaming generator cleanup.""" @pytest.mark.asyncio async def test_streaming_generator_cleanup(self): """ CONCURRENT: Streaming generator cleaned up on abort. Tests that async generators can be properly consumed and stopped. Note: Python doesn't always run finally blocks on generator exit. BUG_PATTERN: Generator not closed on early exit. EXPECTED: Generator can be consumed and stopped. """ chunks_received = [0] async def mock_stream(): """Mock LLM streaming response.""" for i in range(100): yield f"chunk_{i}" await asyncio.sleep(0.001) # Consume only first 3 chunks async for chunk in mock_stream(): chunks_received[0] += 1 if chunks_received[0] >= 3: break # Early exit # Verify we got chunks before stopping assert chunks_received[0] == 3, "Should have received 3 chunks" @pytest.mark.asyncio async def test_streaming_exception_cleanup(self): """ CONCURRENT: Streaming generator cleanup on exception. Tests that generators clean up even when exception occurs during iteration. BUG_PATTERN: Exception prevents cleanup. EXPECTED: Generator cleanup runs despite exception. """ cleanup_called = [0] async def failing_stream(): """Stream that raises exception.""" try: yield "chunk_1" yield "chunk_2" raise ValueError("Stream error") finally: cleanup_called[0] += 1 # Try to consume stream (will fail) chunks = [] try: async for chunk in failing_stream(): chunks.append(chunk) except ValueError: pass # Expected # Verify cleanup called assert cleanup_called[0] == 1, "Cleanup should be called" assert len(chunks) == 2, "Should have 2 chunks before error" class TestWebSocketConnectionCleanup: """Test WebSocket connection cleanup.""" @pytest.mark.asyncio async def test_websocket_cleanup_on_close(self): """ CONCURRENT: WebSocket resources released on close. Tests that WebSocket connections are properly cleaned up when closed during streaming. BUG_PATTERN: WebSocket connections not closed. EXPECTED: Connection resources released. """ connections_opened = [0] connections_closed = [0] class MockWebSocket: """Mock WebSocket for testing.""" def __init__(self): connections_opened[0] += 1 self.closed = False async def send(self, data): """Send data.""" if self.closed: raise ConnectionError("WebSocket closed") async def close(self): """Close connection.""" if not self.closed: self.closed = True connections_closed[0] += 1 async def websocket_streaming(ws: MockWebSocket): """Simulate WebSocket streaming.""" try: for i in range(10): await ws.send(f"message_{i}") await asyncio.sleep(0.01) finally: await ws.close() # Create WebSocket and stream ws = MockWebSocket() task = asyncio.create_task(websocket_streaming(ws)) # Let it send a few messages then cancel await asyncio.sleep(0.05) task.cancel() try: await task except asyncio.CancelledError: pass # Verify WebSocket closed assert ws.closed, "WebSocket should be closed" assert connections_closed[0] == 1, "Close should be called" class TestMemoryLeakDetection: """Test memory leak detection patterns.""" @pytest.mark.asyncio async def test_no_resource_leak_after_many_async_operations(self): """ CONCURRENT: Many async operations should not leak resources excessively. Tests resource usage after many concurrent operations. Resource count should remain within reasonable bounds. Note: Python creates many objects during normal operation. We check for excessive leaks (> 500 objects indicates problem). BUG_PATTERN: Memory leak or connection leak after many operations. EXPECTED: Resource count reasonable (< 500 objects increase). """ # Get initial object count gc.collect() initial_objects = len(gc.get_objects()) # Run many async operations operation_count = 20 async def single_operation(op_id: int): """Single async operation.""" # Simulate some work await asyncio.sleep(0.001) # Create temporary objects data = {"id": op_id, "data": list(range(10))} return data # Launch operations tasks = [single_operation(i) for i in range(operation_count)] results = await asyncio.gather(*tasks, return_exceptions=True) # Verify all succeeded assert all(not isinstance(r, Exception) for r in results) # Force garbage collection gc.collect() final_objects = len(gc.get_objects()) # Verify no excessive leak (allow 500 object tolerance for test infrastructure) object_increase = final_objects - initial_objects assert object_increase < 500, \ f"Possible memory leak: {object_increase} objects added (threshold: 500)" @pytest.mark.asyncio async def test_generator_leak_detection(self): """ CONCURRENT: Detect excessive generator accumulation. Tests that creating and consuming generators doesn't leak. Many generators should be garbage collected after use. BUG_PATTERN: Generators not closed leak resources. EXPECTED: Generator count reasonable. """ # Count objects before gc.collect() objects_before = len(gc.get_objects()) # Create and consume many generators async def mock_generator(): for i in range(10): yield i # Consume generators properly for _ in range(10): async for _ in mock_generator(): pass # Force cleanup gc.collect() objects_after = len(gc.get_objects()) # Allow reasonable increase for test infrastructure object_increase = objects_after - objects_before # Generators should be garbage collected assert object_increase < 200, \ f"Too many objects remaining: {object_increase} increase"