init
This commit is contained in:
@@ -0,0 +1,456 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user