diff --git a/digital-avatar-app/backend/main.py b/digital-avatar-app/backend/main.py index bfb230f..a5d1e4d 100644 --- a/digital-avatar-app/backend/main.py +++ b/digital-avatar-app/backend/main.py @@ -4,7 +4,7 @@ from fastapi.middleware.cors import CORSMiddleware import os import logging -from apscheduler.schedulers.background import BackgroundScheduler +from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.interval import IntervalTrigger import redis as redis_lib @@ -23,6 +23,9 @@ from responses import ok logger = logging.getLogger(__name__) +takeover_scheduler = None +takeover_db = None + app = FastAPI(title="会会数字分身 API", version="1.0.0") app.add_middleware( @@ -110,9 +113,14 @@ def seed(): @app.on_event("startup") def on_startup(): + global takeover_scheduler, takeover_db + init_db() seed() + # Release stale resources when startup is invoked again by a reload/test. + stop_takeover_scheduler() + # --- Takeover scheduler --- try: # Initialize Redis (optional) @@ -137,16 +145,42 @@ def on_startup(): # Initialize takeover service from services.takeover_service import TakeoverService - takeover_service = TakeoverService(SessionLocal(), boxim_client, redis_client) + takeover_db = SessionLocal() + takeover_service = TakeoverService(takeover_db, boxim_client, redis_client) - # Start periodic polling job - scheduler = BackgroundScheduler() - scheduler.add_job( + # AsyncIOScheduler awaits the service coroutine instead of dropping it. + takeover_scheduler = AsyncIOScheduler() + takeover_scheduler.add_job( takeover_service.poll_and_process_messages, trigger=IntervalTrigger(seconds=10), id="takeover_message_poll", ) - scheduler.start() + takeover_scheduler.start() logger.info("Takeover message polling scheduler started (interval=10s)") except Exception as e: + stop_takeover_scheduler() logger.warning(f"Failed to initialize takeover scheduler, app will continue without it: {e}") + + +def stop_takeover_scheduler(): + global takeover_scheduler, takeover_db + + if takeover_scheduler is not None: + try: + if takeover_scheduler.running: + takeover_scheduler.shutdown(wait=False) + except Exception as e: + logger.warning(f"Failed to stop takeover scheduler cleanly: {e}") + finally: + takeover_scheduler = None + + if takeover_db is not None: + try: + takeover_db.close() + finally: + takeover_db = None + + +@app.on_event("shutdown") +def on_shutdown(): + stop_takeover_scheduler() diff --git a/digital-avatar-app/backend/tests/test_takeover_scheduler.py b/digital-avatar-app/backend/tests/test_takeover_scheduler.py index b132770..875bd5b 100644 --- a/digital-avatar-app/backend/tests/test_takeover_scheduler.py +++ b/digital-avatar-app/backend/tests/test_takeover_scheduler.py @@ -14,7 +14,7 @@ def test_app_has_startup_event(): @patch("services.takeover_service.TakeoverService") @patch("services.boxim_client.BoxIMClient") @patch("main.redis_lib.from_url") -@patch("main.BackgroundScheduler") +@patch("main.AsyncIOScheduler") def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls): """Verify scheduler is initialized when Redis is available.""" mock_redis = MagicMock() @@ -25,6 +25,7 @@ def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_ mock_boxim_cls.return_value = mock_boxim mock_takeover = MagicMock() + mock_takeover.poll_and_process_messages = AsyncMock() mock_takeover_cls.return_value = mock_takeover with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379"}): @@ -32,19 +33,23 @@ def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_ on_startup() mock_scheduler_class.return_value.add_job.assert_called_once() + scheduled_callable = mock_scheduler_class.return_value.add_job.call_args.args[0] call_kwargs = mock_scheduler_class.return_value.add_job.call_args[1] + assert scheduled_callable is mock_takeover.poll_and_process_messages assert call_kwargs["id"] == "takeover_message_poll" + mock_scheduler_class.return_value.start.assert_called_once_with() @patch("services.takeover_service.TakeoverService") @patch("services.boxim_client.BoxIMClient") -@patch("main.BackgroundScheduler") +@patch("main.AsyncIOScheduler") def test_scheduler_starts_without_redis(mock_scheduler_class, mock_boxim_cls, mock_takeover_cls): """App should start even when REDIS_URL is not set.""" mock_boxim = MagicMock() mock_boxim_cls.return_value = mock_boxim mock_takeover = MagicMock() + mock_takeover.poll_and_process_messages = AsyncMock() mock_takeover_cls.return_value = mock_takeover with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": ""}, clear=False): @@ -57,7 +62,7 @@ def test_scheduler_starts_without_redis(mock_scheduler_class, mock_boxim_cls, mo @patch("services.takeover_service.TakeoverService") @patch("services.boxim_client.BoxIMClient") @patch("main.redis_lib.from_url") -@patch("main.BackgroundScheduler") +@patch("main.AsyncIOScheduler") def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls): """App should start even when Redis ping fails.""" mock_redis_from_url.side_effect = ConnectionError("Connection refused") @@ -66,6 +71,7 @@ def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from mock_boxim_cls.return_value = mock_boxim mock_takeover = MagicMock() + mock_takeover.poll_and_process_messages = AsyncMock() mock_takeover_cls.return_value = mock_takeover with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://badhost:6379"}): @@ -75,7 +81,7 @@ def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from mock_scheduler_class.return_value.add_job.assert_called_once() -@patch("main.BackgroundScheduler") +@patch("main.AsyncIOScheduler") def test_scheduler_fails_gracefully(mock_scheduler_class): """If scheduler init raises, the app should still start (exception caught).""" mock_scheduler_class.side_effect = RuntimeError("Scheduler crash") @@ -87,6 +93,24 @@ def test_scheduler_fails_gracefully(mock_scheduler_class): # No exception should propagate +def test_scheduler_shutdown_releases_resources(): + """Shutdown should stop polling and close its dedicated database session.""" + import main + + mock_scheduler = MagicMock() + mock_scheduler.running = True + mock_db = MagicMock() + main.takeover_scheduler = mock_scheduler + main.takeover_db = mock_db + + main.on_shutdown() + + mock_scheduler.shutdown.assert_called_once_with(wait=False) + mock_db.close.assert_called_once_with() + assert main.takeover_scheduler is None + assert main.takeover_db is None + + # --- poll_and_process_messages ---