457 lines
13 KiB
Python
457 lines
13 KiB
Python
import json
|
|
from datetime import datetime
|
|
|
|
from services.storage.message_storage import MessageStorage
|
|
|
|
|
|
class _FakeCursor:
|
|
def __init__(self, sink):
|
|
self._sink = sink
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
return False
|
|
|
|
def execute(self, sql, params=None):
|
|
self._sink["sql"] = sql
|
|
self._sink["params"] = params
|
|
self._sink.setdefault("calls", []).append((sql, params))
|
|
return self._sink.get("execute_return", 1)
|
|
|
|
def fetchone(self):
|
|
return self._sink.get("fetchone_result")
|
|
|
|
|
|
class _FakeConn:
|
|
def __init__(self, sink):
|
|
self._sink = sink
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
return False
|
|
|
|
def cursor(self, *args, **kwargs):
|
|
return _FakeCursor(self._sink)
|
|
|
|
|
|
class _ScriptedCursor:
|
|
def __init__(self, steps):
|
|
self._steps = steps
|
|
self._index = 0
|
|
self._current = None
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
return False
|
|
|
|
def execute(self, sql, params=None):
|
|
if self._index >= len(self._steps):
|
|
raise AssertionError(f"Unexpected SQL: {sql}")
|
|
self._current = self._steps[self._index]
|
|
self._current["sql"] = sql
|
|
self._current["params"] = params
|
|
self._index += 1
|
|
return self._current.get("execute_return", 1)
|
|
|
|
def fetchone(self):
|
|
return None if self._current is None else self._current.get("fetchone_result")
|
|
|
|
|
|
class _ScriptedConn:
|
|
def __init__(self, steps):
|
|
self._cursor = _ScriptedCursor(steps)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
return False
|
|
|
|
def cursor(self, *args, **kwargs):
|
|
return self._cursor
|
|
|
|
|
|
def test_save_message_writes_messages_dto_shape(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
sink = {}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
ok = storage.save_message(
|
|
conversation_id="cid-1",
|
|
message_id="mid-1",
|
|
query="select",
|
|
answer="answer",
|
|
workflow_type="conversation",
|
|
user="tester",
|
|
sql_query="SELECT 1",
|
|
execution_result={"total": 1},
|
|
metadata={"trace_id": "t-1"},
|
|
)
|
|
|
|
assert ok is True
|
|
assert "INSERT INTO messages" in sink["sql"]
|
|
|
|
params = sink["params"]
|
|
assert params[0] == "mid-1"
|
|
assert params[1] == "tester"
|
|
assert isinstance(params[2], datetime)
|
|
assert params[3] == "tester"
|
|
assert isinstance(params[4], datetime)
|
|
assert params[5] == "tester"
|
|
assert params[6] == "tester"
|
|
assert params[7] == "mid-1"
|
|
assert params[8] == "cid-1"
|
|
assert params[9] == "tester"
|
|
assert params[10] == "select"
|
|
assert params[11] == "answer"
|
|
assert params[12] is None
|
|
assert params[13] is None
|
|
assert isinstance(params[14], int)
|
|
assert isinstance(params[15], int)
|
|
|
|
log_payload = json.loads(params[16])
|
|
assert log_payload["workflow_type"] == "conversation"
|
|
assert log_payload["sql_query"] == "SELECT 1"
|
|
assert log_payload["execution_result"] == {"total": 1}
|
|
assert log_payload["metadata"] == {"trace_id": "t-1"}
|
|
|
|
|
|
def test_save_message_returns_false_when_disabled(monkeypatch):
|
|
monkeypatch.setattr("services.storage.message_storage.Config.get_section", lambda section: {"enabled": "false"})
|
|
storage = MessageStorage()
|
|
|
|
ok = storage.save_message(conversation_id="cid", message_id="mid", query="q")
|
|
|
|
assert ok is False
|
|
|
|
|
|
def test_save_message_uses_explicit_timestamps(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
sink = {}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
ok = storage.save_message(
|
|
conversation_id="cid-explicit",
|
|
message_id="mid-explicit",
|
|
query="select explicit",
|
|
answer="answer",
|
|
workflow_type="conversation",
|
|
user="tester",
|
|
created_at=1234,
|
|
updated_at=5678,
|
|
)
|
|
|
|
assert ok is True
|
|
assert isinstance(sink["params"][2], datetime)
|
|
assert isinstance(sink["params"][4], datetime)
|
|
assert sink["params"][14] == 1234
|
|
assert sink["params"][15] == 5678
|
|
|
|
|
|
def test_save_message_persists_java_style_logs_into_log_data(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
sink = {}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
ok = storage.save_message(
|
|
conversation_id="cid-log",
|
|
message_id="mid-log",
|
|
query="query",
|
|
answer="answer",
|
|
workflow_type="conversation",
|
|
logs=["line1", "line2"],
|
|
)
|
|
|
|
assert ok is True
|
|
log_payload = json.loads(sink["params"][16])
|
|
assert log_payload["data"] == "line1\nline2"
|
|
|
|
|
|
def test_entity_debug_logging_is_silent_when_switch_disabled(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"entity_debug_enabled": "false",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
|
|
printed = []
|
|
monkeypatch.setattr("builtins.print", lambda *args, **kwargs: printed.append(args[0] if args else ""))
|
|
|
|
storage._log_entity_stage("messages", "create", "start", storage._now_ms(), {"message_id": "mid-1"})
|
|
|
|
assert printed == []
|
|
|
|
|
|
def test_entity_debug_logging_emits_stage_payload_when_switch_enabled(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"entity_debug_enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
|
|
printed = []
|
|
monkeypatch.setattr("builtins.print", lambda *args, **kwargs: printed.append(args[0] if args else ""))
|
|
|
|
storage._log_entity_stage("conversations", "get_by_id", "start", storage._now_ms(), {"conversation_id": "cid-1"})
|
|
|
|
assert len(printed) == 1
|
|
payload = json.loads(printed[0])
|
|
assert payload["level"] == "DEBUG"
|
|
assert payload["event"] == "message_storage.conversations.get_by_id.start"
|
|
assert payload["payload"]["conversation_id"] == "cid-1"
|
|
assert payload["payload"]["entity"] == "conversations"
|
|
assert payload["payload"]["action"] == "get_by_id"
|
|
assert payload["payload"]["stage"] == "start"
|
|
|
|
|
|
def test_update_feedback_by_message_id_success(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
sink = {}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
ok = storage.update_feedback_by_message_id(
|
|
message_id="mid-1",
|
|
feedback="dislike",
|
|
feedback_content="结果不准确",
|
|
)
|
|
|
|
assert ok is True
|
|
assert "UPDATE messages" in sink["sql"]
|
|
assert sink["params"][0] == "dislike"
|
|
assert sink["params"][1] == "结果不准确"
|
|
assert isinstance(sink["params"][2], int)
|
|
assert isinstance(sink["params"][3], datetime)
|
|
assert sink["params"][4] == "mid-1"
|
|
|
|
|
|
def test_update_feedback_by_message_id_not_found_returns_false(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
sink = {"execute_return": 0}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
ok = storage.update_feedback_by_message_id(
|
|
message_id="missing-mid",
|
|
feedback="like",
|
|
)
|
|
|
|
assert ok is False
|
|
|
|
|
|
def test_conversation_create_get_update(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"conversation_table": "conversations",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
storage._conversation_schema_checked = True
|
|
|
|
sink = {}
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _FakeConn(sink))
|
|
|
|
created = storage.create_conversation(
|
|
conversation_id="cid-1",
|
|
user="tester",
|
|
name="hello",
|
|
status="normal",
|
|
introduction="intro",
|
|
created_at=1000,
|
|
updated_at=1000,
|
|
)
|
|
assert created is True
|
|
assert "INSERT INTO conversations" in sink["sql"]
|
|
assert sink["params"][0] == "cid-1"
|
|
assert sink["params"][1] == "tester"
|
|
assert isinstance(sink["params"][2], datetime)
|
|
assert sink["params"][3] == "tester"
|
|
assert isinstance(sink["params"][4], datetime)
|
|
assert sink["params"][7] == "cid-1"
|
|
assert sink["params"][8] == "tester"
|
|
assert sink["params"][9] == "hello"
|
|
|
|
sink["fetchone_result"] = {
|
|
"conversation_id": "cid-1",
|
|
"user": "tester",
|
|
"name": "hello",
|
|
"status": "normal",
|
|
"introduction": "intro",
|
|
"created_at": 1000,
|
|
"updated_at": 1000,
|
|
}
|
|
got = storage.get_conversation_by_id("cid-1")
|
|
assert got is not None
|
|
assert got["conversation_id"] == "cid-1"
|
|
assert got["introduction"] == "intro"
|
|
|
|
updated = storage.update_conversation_updated_at("cid-1", 2000)
|
|
assert updated is True
|
|
assert "UPDATE conversations" in sink["sql"]
|
|
assert sink["params"][0] == 2000
|
|
assert isinstance(sink["params"][1], datetime)
|
|
assert sink["params"][2] == "cid-1"
|
|
|
|
|
|
def test_create_conversation_auto_adds_missing_name_column(monkeypatch):
|
|
monkeypatch.setattr(
|
|
"services.storage.message_storage.Config.get_section",
|
|
lambda section: {
|
|
"enabled": "true",
|
|
"host": "127.0.0.1",
|
|
"port": "3306",
|
|
"user": "root",
|
|
"password": "pwd",
|
|
"database": "db",
|
|
"table": "messages",
|
|
"messages_table": "messages",
|
|
"conversation_table": "conversations",
|
|
"connect_timeout": "5",
|
|
},
|
|
)
|
|
storage = MessageStorage()
|
|
storage._inited = True
|
|
|
|
steps = [
|
|
{"fetchone_result": None},
|
|
{},
|
|
{},
|
|
]
|
|
monkeypatch.setattr(storage, "_get_conn", lambda: _ScriptedConn(steps))
|
|
|
|
created = storage.create_conversation(
|
|
conversation_id="cid-compat",
|
|
user="tester",
|
|
name="new name",
|
|
status="normal",
|
|
introduction=None,
|
|
created_at=1000,
|
|
updated_at=1000,
|
|
)
|
|
|
|
assert created is True
|
|
assert "information_schema.columns" in steps[0]["sql"]
|
|
assert steps[0]["params"] == ("db", "conversations", "name")
|
|
assert "ALTER TABLE conversations" in steps[1]["sql"]
|
|
assert "ADD COLUMN name VARCHAR(255)" in steps[1]["sql"]
|
|
assert "INSERT INTO conversations" in steps[2]["sql"]
|
|
assert steps[2]["params"][0] == "cid-compat"
|
|
assert steps[2]["params"][7] == "cid-compat"
|
|
assert steps[2]["params"][9] == "new name"
|