更新

2026-07-28 15:04:17 +08:00
parent ac406a5f99
commit 8f68af1c2c
27 changed files with 3442 additions and 296 deletions
@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+import json
+import logging.handlers
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from sqlalchemy import text
+from sqlalchemy.pool import StaticPool
+
+
+BACKEND_DIR = Path(__file__).resolve().parents[1]
+if str(BACKEND_DIR) not in sys.path:
+    sys.path.insert(0, str(BACKEND_DIR))
+
+from models.db_config import DatabaseConfig, create_database_engine, engine_kwargs_for_url
+from models.db_migrate import migrate_message_logs_table
+from models.models import MessageLog
+from rpa_engine.douyin_im import protocol
+from rpa_engine.douyin_im.static import Live_pb2, Response_pb2
+from utils import system_logger
+from utils.log_limits import (
+    TRUNCATION_MARKER,
+    bound_message_log_content,
+    bound_raw_message_log_content,
+)
+
+
+class LogLimitTests(unittest.TestCase):
+    def tearDown(self):
+        system_logger.clear()
+
+    def test_system_log_caps_persisted_and_console_detail(self):
+        with patch.dict(os.environ, {"KEFU_SYSTEM_LOG_MAX_CHARS": "512"}):
+            with self.assertLogs("douyin_im.system", level="INFO") as captured:
+                entry = system_logger.record("event", "x" * 5000)
+
+        self.assertLessEqual(len(entry["detail"]), 512)
+        self.assertIn(TRUNCATION_MARKER.strip(), entry["detail"])
+        self.assertLess(len(captured.output[0]), 700)
+
+    def test_oversized_media_log_remains_valid_compact_json(self):
+        payload = json.dumps(
+            {
+                "type": "sticker",
+                "url": "https://example.invalid/sticker.webp",
+                "text": "x" * 20000,
+                "unused_blob": "y" * 20000,
+            },
+            ensure_ascii=False,
+        )
+        with patch.dict(os.environ, {"KEFU_MESSAGE_LOG_MAX_CHARS": "4096"}):
+            bounded = bound_message_log_content(payload)
+
+        decoded = json.loads(bounded)
+        self.assertEqual(decoded["type"], "sticker")
+        self.assertEqual(decoded["url"], "https://example.invalid/sticker.webp")
+        self.assertTrue(decoded["_log_truncated"])
+        self.assertNotIn("unused_blob", decoded)
+        self.assertLessEqual(len(bounded), 4096)
+
+    def test_message_model_validator_caps_all_insert_paths(self):
+        with patch.dict(os.environ, {"KEFU_MESSAGE_LOG_MAX_CHARS": "2048"}):
+            row = MessageLog(message_content="m" * 10000, reply_content="r" * 10000)
+
+        self.assertLessEqual(len(row.message_content), 2048)
+        self.assertLessEqual(len(row.reply_content), 2048)
+
+    def test_raw_message_log_is_bounded(self):
+        with patch.dict(os.environ, {"KEFU_RAW_MESSAGE_LOG_MAX_CHARS": "4096"}):
+            bounded = bound_raw_message_log_content("z" * 20000)
+        self.assertLessEqual(len(bounded), 4096)
+        self.assertIn(TRUNCATION_MARKER.strip(), bounded)
+
+
+class SqliteIoTests(unittest.IsolatedAsyncioTestCase):
+    async def test_short_memory_url_uses_one_static_connection(self):
+        kwargs = engine_kwargs_for_url("sqlite+aiosqlite://")
+        self.assertIs(kwargs["poolclass"], StaticPool)
+        self.assertNotIn("pool_size", kwargs)
+
+        engine = create_database_engine(
+            DatabaseConfig(
+                db_type="sqlite",
+                database_url="sqlite+aiosqlite://",
+            )
+        )
+        try:
+            async with engine.begin() as conn:
+                await conn.execute(text("CREATE TABLE memory_probe (id INTEGER)"))
+            async with engine.begin() as conn:
+                await conn.execute(text("INSERT INTO memory_probe VALUES (1)"))
+                count = (
+                    await conn.execute(text("SELECT count(*) FROM memory_probe"))
+                ).scalar_one()
+            self.assertEqual(count, 1)
+        finally:
+            await engine.dispose()
+
+    async def test_file_sqlite_uses_bounded_pool_and_wal_pragmas(self):
+        kwargs = engine_kwargs_for_url("sqlite+aiosqlite:///example.db")
+        self.assertEqual(kwargs["pool_size"], 5)
+        self.assertEqual(kwargs["max_overflow"], 0)
+        self.assertEqual(kwargs["connect_args"]["timeout"], 30.0)
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            db_path = Path(temp_dir) / "io.db"
+            engine = create_database_engine(
+                DatabaseConfig(db_type="sqlite", db_path=str(db_path))
+            )
+            try:
+                async with engine.connect() as conn:
+                    journal_mode = (await conn.execute(text("PRAGMA journal_mode"))).scalar_one()
+                    synchronous = (await conn.execute(text("PRAGMA synchronous"))).scalar_one()
+                    busy_timeout = (await conn.execute(text("PRAGMA busy_timeout"))).scalar_one()
+                self.assertEqual(str(journal_mode).lower(), "wal")
+                self.assertEqual(synchronous, 1)
+                self.assertEqual(busy_timeout, 30000)
+            finally:
+                await engine.dispose()
+
+    async def test_existing_log_tables_receive_composite_indexes(self):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            db_path = Path(temp_dir) / "migration.db"
+            engine = create_database_engine(
+                DatabaseConfig(db_type="sqlite", db_path=str(db_path))
+            )
+            try:
+                async with engine.begin() as conn:
+                    await conn.execute(
+                        text(
+                            "CREATE TABLE message_logs ("
+                            "id INTEGER PRIMARY KEY, account_id INTEGER, "
+                            "created_at DATETIME, sender_avatar TEXT, status VARCHAR(50))"
+                        )
+                    )
+                    await conn.execute(
+                        text(
+                            "CREATE TABLE received_message_logs ("
+                            "id INTEGER PRIMARY KEY, account_id INTEGER, "
+                            "created_at DATETIME)"
+                        )
+                    )
+                    await conn.execute(
+                        text(
+                            "CREATE TABLE system_logs ("
+                            "id INTEGER PRIMARY KEY, account_id INTEGER, "
+                            "created_at DATETIME)"
+                        )
+                    )
+                    await conn.run_sync(migrate_message_logs_table)
+
+                async with engine.connect() as conn:
+                    message_indexes = {
+                        row[1]
+                        for row in (await conn.execute(text("PRAGMA index_list(message_logs)"))).all()
+                    }
+                    received_indexes = {
+                        row[1]
+                        for row in (
+                            await conn.execute(text("PRAGMA index_list(received_message_logs)"))
+                        ).all()
+                    }
+                    system_indexes = {
+                        row[1]
+                        for row in (await conn.execute(text("PRAGMA index_list(system_logs)"))).all()
+                    }
+                    latest_plan = " ".join(
+                        str(row[-1])
+                        for row in (
+                            await conn.execute(
+                                text(
+                                    "EXPLAIN QUERY PLAN SELECT * FROM message_logs "
+                                    "ORDER BY created_at DESC LIMIT 50"
+                                )
+                            )
+                        ).all()
+                    )
+                    status_plan = " ".join(
+                        str(row[-1])
+                        for row in (
+                            await conn.execute(
+                                text(
+                                    "EXPLAIN QUERY PLAN SELECT count(*) FROM message_logs "
+                                    "WHERE status = 'replied'"
+                                )
+                            )
+                        ).all()
+                    )
+                    account_plan = " ".join(
+                        str(row[-1])
+                        for row in (
+                            await conn.execute(
+                                text(
+                                    "EXPLAIN QUERY PLAN SELECT * FROM message_logs "
+                                    "WHERE account_id = 1 ORDER BY created_at DESC LIMIT 50"
+                                )
+                            )
+                        ).all()
+                    )
+                    system_plan = " ".join(
+                        str(row[-1])
+                        for row in (
+                            await conn.execute(
+                                text(
+                                    "EXPLAIN QUERY PLAN SELECT * FROM system_logs "
+                                    "WHERE account_id = 1 ORDER BY created_at DESC LIMIT 50"
+                                )
+                            )
+                        ).all()
+                    )
+                self.assertIn("ix_message_logs_account_created_at", message_indexes)
+                self.assertIn("ix_message_logs_created_at", message_indexes)
+                self.assertIn("ix_message_logs_status_account_id", message_indexes)
+                self.assertIn(
+                    "ix_received_message_logs_account_created_at",
+                    received_indexes,
+                )
+                self.assertIn("ix_system_logs_account_created_at", system_indexes)
+                self.assertIn("ix_message_logs_created_at", latest_plan)
+                self.assertIn("ix_message_logs_status_account_id", status_plan)
+                self.assertIn("ix_message_logs_account_created_at", account_plan)
+                self.assertIn("ix_system_logs_account_created_at", system_plan)
+            finally:
+                await engine.dispose()
+
+
+class WebSocketDebugTests(unittest.TestCase):
+    def _frame(self, *, message_type: int, content: str) -> bytes:
+        response = Response_pb2.Response()
+        message = response.body.new_message_notify.message
+        message.conversation_id = "0:1:200:100"
+        message.server_message_id = 123
+        message.message_type = message_type
+        message.sender = 200
+        message.content = content
+        frame = Live_pb2.PushFrame()
+        frame.payloadType = "pb"
+        frame.payload = response.SerializeToString()
+        return frame.SerializeToString()
+
+    def test_control_frame_is_filtered_before_debug_writer(self):
+        with patch.object(protocol, "_dump_ws_message") as dump:
+            result = protocol.parse_ws_payload(
+                self._frame(message_type=50001, content='{"command_type":6}')
+            )
+        self.assertEqual(result, [])
+        dump.assert_not_called()
+
+    def test_debug_writer_uses_non_blocking_rotating_queue(self):
+        # Inspect construction without writing chat data to the repository.
+        with tempfile.TemporaryDirectory() as temp_dir:
+            old_path = protocol._WS_DEBUG_PATH
+            protocol._WS_DEBUG_PATH = str(Path(temp_dir) / "ws.log")
+            protocol._WS_DEBUG_LOGGER = None
+            try:
+                with patch.dict(os.environ, {"KEFU_WS_DEBUG": "1"}):
+                    protocol._dump_ws_message(1, "conv", "hello")
+                    logger = protocol._WS_DEBUG_LOGGER
+                self.assertIsNotNone(logger)
+                self.assertIsInstance(logger.handlers[0], logging.handlers.QueueHandler)
+                self.assertIsInstance(
+                    logger._kefu_rotating_handler,
+                    logging.handlers.RotatingFileHandler,
+                )
+            finally:
+                logger = protocol._WS_DEBUG_LOGGER
+                if logger is not None:
+                    logger._kefu_queue_listener.stop()
+                    logger._kefu_rotating_handler.close()
+                    logger.handlers.clear()
+                protocol._WS_DEBUG_LOGGER = None
+                protocol._WS_DEBUG_PATH = old_path
+
+
+if __name__ == "__main__":
+    unittest.main()