MaverickMCP

""" Test graceful shutdown functionality. """ import asyncio import os import signal import subprocess import sys import time from unittest.mock import patch import pytest from maverick_mcp.utils.shutdown import GracefulShutdownHandler, get_shutdown_handler class TestGracefulShutdown: """Test graceful shutdown handler.""" def test_shutdown_handler_creation(self): """Test creating shutdown handler.""" handler = GracefulShutdownHandler("test", shutdown_timeout=10, drain_timeout=5) assert handler.name == "test" assert handler.shutdown_timeout == 10 assert handler.drain_timeout == 5 assert not handler.is_shutting_down() def test_cleanup_registration(self): """Test registering cleanup callbacks.""" handler = GracefulShutdownHandler("test") # Register callbacks callback1_called = False callback2_called = False def callback1(): nonlocal callback1_called callback1_called = True def callback2(): nonlocal callback2_called callback2_called = True handler.register_cleanup(callback1) handler.register_cleanup(callback2) assert len(handler._cleanup_callbacks) == 2 assert callback1 in handler._cleanup_callbacks assert callback2 in handler._cleanup_callbacks @pytest.mark.asyncio async def test_request_tracking(self): """Test request tracking.""" handler = GracefulShutdownHandler("test") # Create mock tasks async def dummy_task(): await asyncio.sleep(0.1) task1 = asyncio.create_task(dummy_task()) task2 = asyncio.create_task(dummy_task()) # Track tasks handler.track_request(task1) handler.track_request(task2) assert len(handler._active_requests) == 2 # Wait for tasks to complete await task1 await task2 await asyncio.sleep(0.1) # Allow cleanup assert len(handler._active_requests) == 0 def test_signal_handler_installation(self): """Test signal handler installation.""" handler = GracefulShutdownHandler("test") # Store original handlers original_sigterm = signal.signal(signal.SIGTERM, signal.SIG_DFL) original_sigint = signal.signal(signal.SIGINT, signal.SIG_DFL) try: # Install handlers handler.install_signal_handlers() # Verify handlers were changed current_sigterm = signal.signal(signal.SIGTERM, signal.SIG_DFL) current_sigint = signal.signal(signal.SIGINT, signal.SIG_DFL) assert current_sigterm == handler._signal_handler assert current_sigint == handler._signal_handler finally: # Restore original handlers signal.signal(signal.SIGTERM, original_sigterm) signal.signal(signal.SIGINT, original_sigint) @pytest.mark.asyncio async def test_async_shutdown_sequence(self): """Test async shutdown sequence.""" handler = GracefulShutdownHandler("test", drain_timeout=0.5) # Track cleanup calls sync_called = False async_called = False def sync_cleanup(): nonlocal sync_called sync_called = True async def async_cleanup(): nonlocal async_called async_called = True handler.register_cleanup(sync_cleanup) handler.register_cleanup(async_cleanup) # Mock sys.exit to prevent actual exit with patch("sys.exit") as mock_exit: # Trigger shutdown handler._shutdown_in_progress = False await handler._async_shutdown("SIGTERM") # Verify shutdown sequence assert handler._shutdown_event.is_set() assert sync_called assert async_called mock_exit.assert_called_once_with(0) @pytest.mark.asyncio async def test_request_draining_timeout(self): """Test request draining with timeout.""" handler = GracefulShutdownHandler("test", drain_timeout=0.2) # Create long-running task async def long_task(): await asyncio.sleep(1.0) # Longer than drain timeout task = asyncio.create_task(long_task()) handler.track_request(task) # Start draining start_time = time.time() try: await asyncio.wait_for(handler._wait_for_requests(), timeout=0.3) except TimeoutError: pass drain_time = time.time() - start_time # Should timeout quickly since task won't complete assert drain_time < 0.5 assert task in handler._active_requests # Cancel task to clean up task.cancel() try: await task except asyncio.CancelledError: pass def test_global_shutdown_handler(self): """Test global shutdown handler singleton.""" handler1 = get_shutdown_handler("test1") handler2 = get_shutdown_handler("test2") # Should return same instance assert handler1 is handler2 assert handler1.name == "test1" # First call sets the name @pytest.mark.asyncio async def test_cleanup_callback_error_handling(self): """Test error handling in cleanup callbacks.""" handler = GracefulShutdownHandler("test") # Create callback that raises exception def failing_callback(): raise RuntimeError("Cleanup failed") async def async_failing_callback(): raise RuntimeError("Async cleanup failed") handler.register_cleanup(failing_callback) handler.register_cleanup(async_failing_callback) # Mock sys.exit with patch("sys.exit"): # Should not raise despite callback errors await handler._async_shutdown("SIGTERM") # Handler should still complete shutdown assert handler._shutdown_event.is_set() @pytest.mark.asyncio async def test_sync_request_tracking(self): """Test synchronous request tracking context manager.""" handler = GracefulShutdownHandler("test") # Use context manager with handler.track_sync_request(): # In real usage, this would track the request pass # Should complete without error assert True @pytest.mark.skipif( sys.platform == "win32", reason="SIGHUP not available on Windows" ) def test_sighup_handling(self): """Test SIGHUP signal handling.""" handler = GracefulShutdownHandler("test") # Store original handler original_sighup = signal.signal(signal.SIGHUP, signal.SIG_DFL) try: handler.install_signal_handlers() # Verify SIGHUP handler was installed current_sighup = signal.signal(signal.SIGHUP, signal.SIG_DFL) assert current_sighup == handler._signal_handler finally: # Restore original handler signal.signal(signal.SIGHUP, original_sighup) @pytest.mark.integration class TestGracefulShutdownIntegration: """Integration tests for graceful shutdown.""" @pytest.mark.asyncio async def test_server_graceful_shutdown(self): """Test actual server graceful shutdown.""" # This would test with a real server process # For now, we'll simulate it # Start a subprocess that uses our shutdown handler script = """ import asyncio import signal import sys import time import os sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from maverick_mcp.utils.shutdown import graceful_shutdown async def main(): with graceful_shutdown("test-server") as handler: # Simulate server running print("Server started", flush=True) # Wait for shutdown try: await handler.wait_for_shutdown() except KeyboardInterrupt: pass print("Server shutting down", flush=True) if __name__ == "__main__": asyncio.run(main()) """ # Write script to temp file import tempfile with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: f.write(script) script_path = f.name try: # Start subprocess proc = subprocess.Popen( [sys.executable, script_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) # Wait for startup await asyncio.sleep(0.5) # Send SIGTERM proc.send_signal(signal.SIGTERM) # Wait for completion stdout, stderr = proc.communicate(timeout=5) # Verify graceful shutdown assert "Server started" in stdout assert "Server shutting down" in stdout assert proc.returncode == 0 finally: os.unlink(script_path)