feat: add scheduled message polling for avatar takeover
Add APScheduler BackgroundScheduler to on_startup that polls takeover_service.poll_and_process_messages every 10 seconds. Added redis and apscheduler dependencies. Added poll_and_process_messages, fetch_unread_messages, and process_message methods to TakeoverService. Scheduler initialization is graceful — Redis/Box IM failures do not prevent app startup.
This commit is contained in:
@@ -2,6 +2,11 @@ from fastapi import FastAPI
|
|||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
import os
|
import os
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from apscheduler.schedulers.background import BackgroundScheduler
|
||||||
|
from apscheduler.triggers.interval import IntervalTrigger
|
||||||
|
import redis as redis_lib
|
||||||
|
|
||||||
from database import init_db, SessionLocal
|
from database import init_db, SessionLocal
|
||||||
from models import Avatar, Authorization, Organization, TokenAccount, TokenPlan
|
from models import Avatar, Authorization, Organization, TokenAccount, TokenPlan
|
||||||
@@ -15,6 +20,8 @@ import routers.huihui_auth
|
|||||||
import routers.chat
|
import routers.chat
|
||||||
from responses import ok
|
from responses import ok
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
app = FastAPI(title="会会数字分身 API", version="1.0.0")
|
app = FastAPI(title="会会数字分身 API", version="1.0.0")
|
||||||
|
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
@@ -103,3 +110,41 @@ def seed():
|
|||||||
def on_startup():
|
def on_startup():
|
||||||
init_db()
|
init_db()
|
||||||
seed()
|
seed()
|
||||||
|
|
||||||
|
# --- Takeover scheduler ---
|
||||||
|
try:
|
||||||
|
# Initialize Redis (optional)
|
||||||
|
redis_client = None
|
||||||
|
redis_url = os.getenv("REDIS_URL", "")
|
||||||
|
if redis_url:
|
||||||
|
try:
|
||||||
|
redis_client = redis_lib.from_url(redis_url)
|
||||||
|
redis_client.ping()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Redis connection failed, delayed takeover will degrade to immediate: {e}")
|
||||||
|
|
||||||
|
# Initialize Box IM client
|
||||||
|
from services.boxim_client import BoxIMClient
|
||||||
|
boxim_config = {
|
||||||
|
"HUIHUI_IM_BASE_URL": os.getenv("HUIHUI_IM_BASE_URL", "http://192.168.1.200:60040"),
|
||||||
|
"HUIHUI_APP_ID": os.getenv("HUIHUI_APP_ID", ""),
|
||||||
|
"HUIHUI_ACCESS_ID": os.getenv("HUIHUI_ACCESS_ID", ""),
|
||||||
|
"HUIHUI_ACCESS_SECRET": os.getenv("HUIHUI_ACCESS_SECRET", ""),
|
||||||
|
}
|
||||||
|
boxim_client = BoxIMClient(boxim_config)
|
||||||
|
|
||||||
|
# Initialize takeover service
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
takeover_service = TakeoverService(SessionLocal(), boxim_client, redis_client)
|
||||||
|
|
||||||
|
# Start periodic polling job
|
||||||
|
scheduler = BackgroundScheduler()
|
||||||
|
scheduler.add_job(
|
||||||
|
takeover_service.poll_and_process_messages,
|
||||||
|
trigger=IntervalTrigger(seconds=10),
|
||||||
|
id="takeover_message_poll",
|
||||||
|
)
|
||||||
|
scheduler.start()
|
||||||
|
logger.info("Takeover message polling scheduler started (interval=10s)")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to initialize takeover scheduler, app will continue without it: {e}")
|
||||||
|
|||||||
@@ -7,3 +7,5 @@ httpx
|
|||||||
pypdf
|
pypdf
|
||||||
python-docx
|
python-docx
|
||||||
openpyxl
|
openpyxl
|
||||||
|
redis>=5.0
|
||||||
|
apscheduler>=3.10
|
||||||
|
|||||||
@@ -144,3 +144,31 @@ class TakeoverService:
|
|||||||
self.redis.delete(key)
|
self.redis.delete(key)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to process delayed queue: {e}")
|
logger.error(f"Failed to process delayed queue: {e}")
|
||||||
|
|
||||||
|
async def poll_and_process_messages(self):
|
||||||
|
"""Periodic polling job: fetch unread messages and process each."""
|
||||||
|
try:
|
||||||
|
messages = await self.fetch_unread_messages()
|
||||||
|
for msg in messages:
|
||||||
|
await self.process_message(msg)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"poll_and_process_messages failed: {e}")
|
||||||
|
|
||||||
|
async def fetch_unread_messages(self) -> list:
|
||||||
|
"""Fetch unread messages from Box IM. Stub — replace with real API call."""
|
||||||
|
logger.debug("fetch_unread_messages: no real API wired yet")
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def process_message(self, message: dict):
|
||||||
|
"""Process a single message: check takeover, dispatch immediate or delayed."""
|
||||||
|
owner_id = message.get("owner_huihui_id", "")
|
||||||
|
from_id = message.get("from_accid", "")
|
||||||
|
|
||||||
|
auth = self.check_takeover_enabled(owner_id, from_id)
|
||||||
|
if not auth:
|
||||||
|
return
|
||||||
|
|
||||||
|
if auth.takeover_mode == "immediate":
|
||||||
|
await self.execute_takeover(auth, message)
|
||||||
|
else:
|
||||||
|
self.enqueue_delayed_message(auth, message)
|
||||||
|
|||||||
@@ -0,0 +1,197 @@
|
|||||||
|
"""Tests for the scheduled takeover message polling."""
|
||||||
|
import json
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import MagicMock, patch, AsyncMock
|
||||||
|
|
||||||
|
|
||||||
|
def test_app_has_startup_event():
|
||||||
|
"""Verify the app has a startup event configured."""
|
||||||
|
from main import app
|
||||||
|
startup_handlers = [handler for handler in app.router.on_startup]
|
||||||
|
assert len(startup_handlers) > 0
|
||||||
|
|
||||||
|
|
||||||
|
@patch("services.takeover_service.TakeoverService")
|
||||||
|
@patch("services.boxim_client.BoxIMClient")
|
||||||
|
@patch("main.redis_lib.from_url")
|
||||||
|
@patch("main.BackgroundScheduler")
|
||||||
|
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()
|
||||||
|
mock_redis.ping.return_value = None
|
||||||
|
mock_redis_from_url.return_value = mock_redis
|
||||||
|
|
||||||
|
mock_boxim = MagicMock()
|
||||||
|
mock_boxim_cls.return_value = mock_boxim
|
||||||
|
|
||||||
|
mock_takeover = MagicMock()
|
||||||
|
mock_takeover_cls.return_value = mock_takeover
|
||||||
|
|
||||||
|
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379"}):
|
||||||
|
from main import on_startup
|
||||||
|
on_startup()
|
||||||
|
|
||||||
|
mock_scheduler_class.return_value.add_job.assert_called_once()
|
||||||
|
call_kwargs = mock_scheduler_class.return_value.add_job.call_args[1]
|
||||||
|
assert call_kwargs["id"] == "takeover_message_poll"
|
||||||
|
|
||||||
|
|
||||||
|
@patch("services.takeover_service.TakeoverService")
|
||||||
|
@patch("services.boxim_client.BoxIMClient")
|
||||||
|
@patch("main.BackgroundScheduler")
|
||||||
|
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_cls.return_value = mock_takeover
|
||||||
|
|
||||||
|
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": ""}, clear=False):
|
||||||
|
from main import on_startup
|
||||||
|
on_startup()
|
||||||
|
|
||||||
|
mock_scheduler_class.return_value.add_job.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@patch("services.takeover_service.TakeoverService")
|
||||||
|
@patch("services.boxim_client.BoxIMClient")
|
||||||
|
@patch("main.redis_lib.from_url")
|
||||||
|
@patch("main.BackgroundScheduler")
|
||||||
|
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")
|
||||||
|
|
||||||
|
mock_boxim = MagicMock()
|
||||||
|
mock_boxim_cls.return_value = mock_boxim
|
||||||
|
|
||||||
|
mock_takeover = MagicMock()
|
||||||
|
mock_takeover_cls.return_value = mock_takeover
|
||||||
|
|
||||||
|
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://badhost:6379"}):
|
||||||
|
from main import on_startup
|
||||||
|
on_startup()
|
||||||
|
|
||||||
|
mock_scheduler_class.return_value.add_job.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@patch("main.BackgroundScheduler")
|
||||||
|
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")
|
||||||
|
|
||||||
|
with patch("main.init_db"), patch("main.seed"):
|
||||||
|
from main import on_startup
|
||||||
|
on_startup()
|
||||||
|
|
||||||
|
# No exception should propagate
|
||||||
|
|
||||||
|
|
||||||
|
# --- poll_and_process_messages ---
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_db():
|
||||||
|
return MagicMock()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_boxim():
|
||||||
|
return AsyncMock()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_poll_and_process_messages_calls_fetch_and_process(mock_db, mock_boxim):
|
||||||
|
"""poll_and_process_messages should fetch messages and process each."""
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
|
||||||
|
service = TakeoverService(mock_db, mock_boxim)
|
||||||
|
service.fetch_unread_messages = AsyncMock(return_value=[
|
||||||
|
{"owner_huihui_id": "owner_1", "from_accid": "user_1", "content": "hi"},
|
||||||
|
{"owner_huihui_id": "owner_2", "from_accid": "user_2", "content": "hello"},
|
||||||
|
])
|
||||||
|
service.process_message = AsyncMock()
|
||||||
|
|
||||||
|
await service.poll_and_process_messages()
|
||||||
|
|
||||||
|
service.fetch_unread_messages.assert_awaited_once()
|
||||||
|
assert service.process_message.await_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_poll_and_process_messages_handles_errors(mock_db, mock_boxim):
|
||||||
|
"""poll_and_process_messages should not crash on fetch failure."""
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
|
||||||
|
service = TakeoverService(mock_db, mock_boxim)
|
||||||
|
service.fetch_unread_messages = AsyncMock(side_effect=ConnectionError("Box IM down"))
|
||||||
|
|
||||||
|
await service.poll_and_process_messages()
|
||||||
|
# No exception should propagate
|
||||||
|
|
||||||
|
|
||||||
|
# --- process_message ---
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_auth():
|
||||||
|
auth = MagicMock()
|
||||||
|
auth.takeover_enabled = True
|
||||||
|
auth.takeover_mode = "immediate"
|
||||||
|
auth.takeover_delay_seconds = 30
|
||||||
|
auth.avatar_id = "avatar_123"
|
||||||
|
auth.target_id = "target_user_123"
|
||||||
|
return auth
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_process_message_immediate_mode(mock_db, mock_boxim, mock_auth):
|
||||||
|
"""When takeover_mode is 'immediate', execute_takeover should be called."""
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
|
||||||
|
service = TakeoverService(mock_db, mock_boxim)
|
||||||
|
service.check_takeover_enabled = MagicMock(return_value=mock_auth)
|
||||||
|
service.execute_takeover = AsyncMock(return_value=True)
|
||||||
|
service.enqueue_delayed_message = MagicMock()
|
||||||
|
|
||||||
|
message = {"owner_huihui_id": "owner_1", "from_accid": "user_1", "content": "hi"}
|
||||||
|
await service.process_message(message)
|
||||||
|
|
||||||
|
service.execute_takeover.assert_awaited_once_with(mock_auth, message)
|
||||||
|
service.enqueue_delayed_message.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_process_message_delayed_mode(mock_db, mock_boxim, mock_auth):
|
||||||
|
"""When takeover_mode is not 'immediate', message should be enqueued."""
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
|
||||||
|
mock_auth.takeover_mode = "delayed"
|
||||||
|
|
||||||
|
service = TakeoverService(mock_db, mock_boxim)
|
||||||
|
service.check_takeover_enabled = MagicMock(return_value=mock_auth)
|
||||||
|
service.execute_takeover = AsyncMock()
|
||||||
|
service.enqueue_delayed_message = MagicMock()
|
||||||
|
|
||||||
|
message = {"owner_huihui_id": "owner_1", "from_accid": "user_1", "content": "hi"}
|
||||||
|
await service.process_message(message)
|
||||||
|
|
||||||
|
service.enqueue_delayed_message.assert_called_once_with(mock_auth, message)
|
||||||
|
service.execute_takeover.assert_not_awaited()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_process_message_no_takeover(mock_db, mock_boxim):
|
||||||
|
"""When takeover is not enabled, nothing should happen."""
|
||||||
|
from services.takeover_service import TakeoverService
|
||||||
|
|
||||||
|
service = TakeoverService(mock_db, mock_boxim)
|
||||||
|
service.check_takeover_enabled = MagicMock(return_value=None)
|
||||||
|
service.execute_takeover = AsyncMock()
|
||||||
|
service.enqueue_delayed_message = MagicMock()
|
||||||
|
|
||||||
|
message = {"owner_huihui_id": "owner_1", "from_accid": "user_1", "content": "hi"}
|
||||||
|
await service.process_message(message)
|
||||||
|
|
||||||
|
service.execute_takeover.assert_not_awaited()
|
||||||
|
service.enqueue_delayed_message.assert_not_called()
|
||||||
Reference in New Issue
Block a user