feat(memory): add memory gateway client and service
This commit is contained in:
242
app-instance/backend/tests/unit/test_memory_gateway_service.py
Normal file
242
app-instance/backend/tests/unit/test_memory_gateway_service.py
Normal file
@ -0,0 +1,242 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from beaver.foundation.config import MemoryGatewayConfig
|
||||
from beaver.integrations.memory_gateway import MemoryGatewayClient, MemoryGatewayClientError
|
||||
from beaver.services.memory_gateway_service import MemoryGatewayService
|
||||
|
||||
|
||||
def _config() -> MemoryGatewayConfig:
|
||||
return MemoryGatewayConfig(
|
||||
base_url="http://gateway.test",
|
||||
user_id="gateway-user",
|
||||
user_key="uk_super_secret",
|
||||
app_id="beaver",
|
||||
project_id="sandbox",
|
||||
scope=["current_chat", "resources"],
|
||||
top_k=5,
|
||||
timeout_seconds=7.5,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_uses_exact_gateway_paths_and_payloads() -> None:
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if request.url.path == "/memories/search":
|
||||
return httpx.Response(200, json={"results": []})
|
||||
return httpx.Response(200, json={"session_id": "chat:web:alpha", "backend": {"data": {"status": "ok"}}})
|
||||
|
||||
client = MemoryGatewayClient(_config(), transport=httpx.MockTransport(handler))
|
||||
|
||||
await client.search({"query": "hello"})
|
||||
await client.add({"session_id": "chat:web:alpha", "messages": []})
|
||||
await client.flush({"session_id": "chat:web:alpha"})
|
||||
|
||||
assert [request.url.path for request in requests] == [
|
||||
"/memories/search",
|
||||
"/memories/add",
|
||||
"/memories/flush",
|
||||
]
|
||||
assert [json.loads(request.content) for request in requests] == [
|
||||
{"query": "hello"},
|
||||
{"session_id": "chat:web:alpha", "messages": []},
|
||||
{"session_id": "chat:web:alpha"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_error_is_sanitized() -> None:
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(401, json={"detail": "uk_super_secret rejected"})
|
||||
|
||||
client = MemoryGatewayClient(_config(), transport=httpx.MockTransport(handler))
|
||||
|
||||
with pytest.raises(MemoryGatewayClientError) as exc_info:
|
||||
await client.search({"user_key": "uk_super_secret"})
|
||||
|
||||
assert exc_info.value.operation == "search"
|
||||
assert exc_info.value.status_code == 401
|
||||
assert "uk_super_secret" not in str(exc_info.value)
|
||||
|
||||
|
||||
class FakeGatewayClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
search_response: dict | None = None,
|
||||
add_error: MemoryGatewayClientError | None = None,
|
||||
flush_error: MemoryGatewayClientError | None = None,
|
||||
) -> None:
|
||||
self.search_response = search_response or {"results": []}
|
||||
self.add_error = add_error
|
||||
self.flush_error = flush_error
|
||||
self.calls: list[tuple[str, dict]] = []
|
||||
|
||||
async def search(self, payload: dict) -> dict:
|
||||
self.calls.append(("search", payload))
|
||||
return self.search_response
|
||||
|
||||
async def add(self, payload: dict) -> dict:
|
||||
self.calls.append(("add", payload))
|
||||
if self.add_error:
|
||||
raise self.add_error
|
||||
return {"session_id": payload["session_id"]}
|
||||
|
||||
async def flush(self, payload: dict) -> dict:
|
||||
self.calls.append(("flush", payload))
|
||||
if self.flush_error:
|
||||
raise self.flush_error
|
||||
return {"session_id": payload["session_id"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_sanitizes_results_and_builds_reference_message() -> None:
|
||||
client = FakeGatewayClient(
|
||||
search_response={
|
||||
"results": [
|
||||
{
|
||||
"id": "mem-1",
|
||||
"session_id": "chat:web:alpha",
|
||||
"text": "The user uploaded a contract.",
|
||||
"score": 0.91,
|
||||
"source_scope": "resources",
|
||||
"resource_uri": "resource://gateway-user/r1",
|
||||
"raw": {"secret_backend_detail": "discard-me"},
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
service = MemoryGatewayService(_config(), client=client)
|
||||
|
||||
outcome = await service.recall_before_run(session_id="web:alpha", query="contract")
|
||||
|
||||
assert outcome.error is None
|
||||
assert outcome.result_count == 1
|
||||
assert client.calls == [
|
||||
(
|
||||
"search",
|
||||
{
|
||||
"user_id": "gateway-user",
|
||||
"user_key": "uk_super_secret",
|
||||
"conversation_id": "web:alpha",
|
||||
"query": "contract",
|
||||
"scope": ["current_chat", "resources"],
|
||||
"top_k": 5,
|
||||
"app_id": "beaver",
|
||||
"project_id": "sandbox",
|
||||
},
|
||||
)
|
||||
]
|
||||
assert len(outcome.reference_messages) == 1
|
||||
message = outcome.reference_messages[0]
|
||||
assert message["role"] == "user"
|
||||
assert "The user uploaded a contract." in message["content"]
|
||||
assert "discard-me" not in message["content"]
|
||||
assert "untrusted reference data" in message["content"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_rejects_malformed_results_shape() -> None:
|
||||
service = MemoryGatewayService(
|
||||
_config(),
|
||||
client=FakeGatewayClient(search_response={"results": {"not": "a list"}}),
|
||||
)
|
||||
|
||||
outcome = await service.recall_before_run(session_id="web:alpha", query="contract")
|
||||
|
||||
assert outcome.reference_messages == []
|
||||
assert outcome.result_count == 0
|
||||
assert outcome.error is not None
|
||||
assert outcome.error.category == "invalid_response"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_after_run_adds_two_messages_then_flushes() -> None:
|
||||
client = FakeGatewayClient()
|
||||
service = MemoryGatewayService(_config(), client=client)
|
||||
|
||||
outcome = await service.persist_after_run(
|
||||
session_id="web:alpha",
|
||||
user_text="hello",
|
||||
assistant_text="hi",
|
||||
user_timestamp_ms=1000,
|
||||
assistant_timestamp_ms=1001,
|
||||
)
|
||||
|
||||
assert outcome.add_succeeded is True
|
||||
assert outcome.flush_succeeded is True
|
||||
assert outcome.add_error is None
|
||||
assert outcome.flush_error is None
|
||||
assert client.calls == [
|
||||
(
|
||||
"add",
|
||||
{
|
||||
"user_id": "gateway-user",
|
||||
"user_key": "uk_super_secret",
|
||||
"session_id": "chat:web:alpha",
|
||||
"app_id": "beaver",
|
||||
"project_id": "sandbox",
|
||||
"messages": [
|
||||
{"sender_id": "gateway-user", "role": "user", "timestamp": 1000, "content": "hello"},
|
||||
{"sender_id": "beaver", "role": "assistant", "timestamp": 1001, "content": "hi"},
|
||||
],
|
||||
},
|
||||
),
|
||||
(
|
||||
"flush",
|
||||
{
|
||||
"user_id": "gateway-user",
|
||||
"user_key": "uk_super_secret",
|
||||
"session_id": "chat:web:alpha",
|
||||
"app_id": "beaver",
|
||||
"project_id": "sandbox",
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_failure_skips_flush() -> None:
|
||||
add_error = MemoryGatewayClientError("add", "http_status", status_code=503)
|
||||
client = FakeGatewayClient(add_error=add_error)
|
||||
service = MemoryGatewayService(_config(), client=client)
|
||||
|
||||
outcome = await service.persist_after_run(
|
||||
session_id="web:alpha",
|
||||
user_text="hello",
|
||||
assistant_text="hi",
|
||||
user_timestamp_ms=1000,
|
||||
assistant_timestamp_ms=1001,
|
||||
)
|
||||
|
||||
assert outcome.add_succeeded is False
|
||||
assert outcome.flush_succeeded is False
|
||||
assert outcome.add_error is add_error
|
||||
assert [name for name, _ in client.calls] == ["add"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_failure_preserves_successful_add() -> None:
|
||||
flush_error = MemoryGatewayClientError("flush", "network")
|
||||
client = FakeGatewayClient(flush_error=flush_error)
|
||||
service = MemoryGatewayService(_config(), client=client)
|
||||
|
||||
outcome = await service.persist_after_run(
|
||||
session_id="web:alpha",
|
||||
user_text="hello",
|
||||
assistant_text="hi",
|
||||
user_timestamp_ms=1000,
|
||||
assistant_timestamp_ms=1001,
|
||||
)
|
||||
|
||||
assert outcome.add_succeeded is True
|
||||
assert outcome.flush_succeeded is False
|
||||
assert outcome.flush_error is flush_error
|
||||
assert [name for name, _ in client.calls] == ["add", "flush"]
|
||||
Reference in New Issue
Block a user