282 lines
12 KiB
Python
282 lines
12 KiB
Python
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()
|