126 lines
3.6 KiB
Python
126 lines
3.6 KiB
Python
import asyncio
|
|
|
|
from memory_system_api.clients import EverOSMemorySystemClient, OpenVikingMemorySystemClient
|
|
|
|
|
|
class FakeStore:
|
|
def __init__(self):
|
|
self.saved = {}
|
|
|
|
def get_user_key(self, user_id: str) -> str | None:
|
|
return self.saved.get(user_id)
|
|
|
|
def save_user_key(self, user_id: str, user_key: str) -> None:
|
|
self.saved[user_id] = user_key
|
|
|
|
|
|
class FakeResponse:
|
|
def __init__(self, status_code: int, data: dict):
|
|
self.status_code = status_code
|
|
self._data = data
|
|
|
|
def json(self) -> dict:
|
|
return self._data
|
|
|
|
def raise_for_status(self) -> None:
|
|
if self.status_code >= 400:
|
|
raise AssertionError(f"unexpected status {self.status_code}")
|
|
|
|
|
|
class FakeAsyncClient:
|
|
def __init__(self, calls: list, responses: list[FakeResponse], api_key: str, headers: dict):
|
|
self.calls = calls
|
|
self.responses = responses
|
|
self.api_key = api_key
|
|
self.headers = headers
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
async def post(self, path: str, json: dict | None = None) -> FakeResponse:
|
|
self.calls.append(("post", self.api_key, self.headers, path, json))
|
|
return self.responses.pop(0)
|
|
|
|
|
|
def test_openviking_uses_root_identity_when_account_already_exists():
|
|
store = FakeStore()
|
|
client = OpenVikingMemorySystemClient(store=store)
|
|
client.root_key = "root-key"
|
|
calls = []
|
|
responses = [FakeResponse(409, {"status": "error", "error": {"code": "CONFLICT"}})]
|
|
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
|
|
credential = asyncio.run(client.ensure_user("tom"))
|
|
|
|
assert credential.api_key == "root-key"
|
|
assert credential.account_id == "tom"
|
|
assert credential.user_id == "tom"
|
|
assert store.saved == {}
|
|
|
|
|
|
def test_openviking_root_identity_headers_are_sent_for_session_create():
|
|
client = OpenVikingMemorySystemClient(store=FakeStore())
|
|
client.root_key = "root-key"
|
|
calls = []
|
|
responses = [FakeResponse(200, {"status": "ok", "result": {"session_id": "sess-2"}})]
|
|
client._client = lambda api_key, extra_headers=None: FakeAsyncClient( # type: ignore[method-assign]
|
|
calls,
|
|
responses,
|
|
api_key,
|
|
extra_headers or {},
|
|
)
|
|
credential = client.root_credential("tom")
|
|
|
|
result = asyncio.run(client.ensure_session(credential, "sess-2"))
|
|
|
|
assert result == {"status": "ok", "result": {"session_id": "sess-2"}}
|
|
assert calls == [
|
|
(
|
|
"post",
|
|
"root-key",
|
|
{"X-OpenViking-Account": "tom", "X-OpenViking-User": "tom"},
|
|
"/api/v1/sessions",
|
|
{"session_id": "sess-2"},
|
|
)
|
|
]
|
|
|
|
|
|
def test_everos_assistant_payload_does_not_use_user_id_as_sender():
|
|
client = EverOSMemorySystemClient()
|
|
|
|
payload = client.build_message_payload(
|
|
user_id="tom",
|
|
session_id="sess-1",
|
|
role="assistant",
|
|
content="我记住了",
|
|
)
|
|
|
|
message = payload["messages"][0]
|
|
assert message["role"] == "assistant"
|
|
assert message["sender_id"] != "tom"
|
|
assert message["sender_name"] != "tom"
|
|
|
|
|
|
def test_everos_user_payload_uses_user_id_as_sender():
|
|
client = EverOSMemorySystemClient()
|
|
|
|
payload = client.build_message_payload(
|
|
user_id="tom",
|
|
session_id="sess-1",
|
|
role="user",
|
|
content="我喜欢拿铁",
|
|
)
|
|
|
|
message = payload["messages"][0]
|
|
assert message["role"] == "user"
|
|
assert message["sender_id"] == "tom"
|
|
assert message["sender_name"] == "tom"
|