backend/tests/test_memory_api.py - master - SupremeSoviet/llm-memory

backend/tests/test_memory_api.py
162 строки6 KB
Vladimir Gubin
feat: add local memory dialogue system
25 май 2026, 01:05
25 май 2026, 01:050f5c6d4
from __future__ import annotations

from datetime import UTC, datetime

import pytest
from httpx import ASGITransport, AsyncClient

from app.main import create_application
from app.memory.storage import MemoryStorage
from tests.fakes import FakeEmbedder, FakeQdrantClient


@pytest.fixture
async def memory_test_context() -> tuple[AsyncClient, FakeDatabase, FakeQdrantClient, FakeEmbedder]:
    database = FakeDatabase()
    application = create_application()
    qdrant_client = FakeQdrantClient()
    embedder = FakeEmbedder(
        {
            ("любит пиццу", "passage"): [1.0, 0.0],
            ("живет в Челябинске", "passage"): [0.0, 1.0],
        }
    )
    application.state.mongodb_database = database
    application.state.qdrant_client = qdrant_client
    application.state.embedder = embedder
    transport = ASGITransport(app=application)
    async with AsyncClient(transport=transport, base_url="http://test") as client:
        yield client, database, qdrant_client, embedder


class FakeCollection:
    def __init__(self) -> None:
        self.documents: list[dict] = []

    async def find_one(self, query: dict) -> dict | None:
        session_identifier = query["session_id"]
        for document in self.documents:
            if document["session_id"] == session_identifier:
                return dict(document)
        return None


class FakeDatabase:
    def __init__(self) -> None:
        self.collections = {"dialogue_sessions": FakeCollection()}

    def __getitem__(self, collection_name: str) -> FakeCollection:
        return self.collections[collection_name]


@pytest.mark.asyncio
async def test_list_memory_records_returns_memories_for_session(
    memory_test_context: tuple[AsyncClient, FakeDatabase, FakeQdrantClient, FakeEmbedder],
) -> None:
    memory_test_client, database, qdrant_client, embedder = memory_test_context
    database["dialogue_sessions"].documents = [
        {
            "session_id": "session-1",
            "persona_id": "persona-1",
            "messages": [],
            "created_at": datetime.now(UTC),
            "updated_at": datetime.now(UTC),
        }
    ]
    storage = MemoryStorage(
        qdrant_client=qdrant_client,
        embedder=embedder,
    )
    await storage.store("persona-1", "favorite_food", "любит пиццу")
    await storage.store("persona-1", "city", "живет в Челябинске")
    await storage.store("persona-2", "pet", "любит пиццу")

    response = await memory_test_client.get("/api/memory/session-1")

    assert response.status_code == 200
    response_payload = response.json()
    assert response_payload["session_id"] == "session-1"
    assert response_payload["persona_id"] == "persona-1"
    assert [memory["key"] for memory in response_payload["memories"]] == ["city", "favorite_food"]


@pytest.mark.asyncio
async def test_clear_memory_records_deletes_only_target_session(
    memory_test_context: tuple[AsyncClient, FakeDatabase, FakeQdrantClient, FakeEmbedder],
) -> None:
    memory_test_client, database, qdrant_client, embedder = memory_test_context
    database["dialogue_sessions"].documents = [
        {
            "session_id": "session-1",
            "persona_id": "persona-1",
            "messages": [],
            "created_at": datetime.now(UTC),
            "updated_at": datetime.now(UTC),
        },
        {
            "session_id": "session-2",
            "persona_id": "persona-2",
            "messages": [],
            "created_at": datetime.now(UTC),
            "updated_at": datetime.now(UTC),
        },
    ]
    storage = MemoryStorage(
        qdrant_client=qdrant_client,
        embedder=embedder,
    )
    await storage.store("persona-1", "favorite_food", "любит пиццу")
    await storage.store("persona-1", "city", "живет в Челябинске")
    await storage.store("persona-2", "pet", "любит пиццу")

    response = await memory_test_client.delete("/api/memory/session-1")

    assert response.status_code == 200
    assert response.json() == {"session_id": "session-1", "deleted_count": 2}

    session_one_response = await memory_test_client.get("/api/memory/session-1")
    session_two_response = await memory_test_client.get("/api/memory/session-2")
    assert session_one_response.json()["memories"] == []
    assert len(session_two_response.json()["memories"]) == 1


@pytest.mark.asyncio
async def test_clear_memory_by_persona_deletes_without_active_session(
    memory_test_context: tuple[AsyncClient, FakeDatabase, FakeQdrantClient, FakeEmbedder],
) -> None:
    memory_test_client, _, qdrant_client, embedder = memory_test_context
    storage = MemoryStorage(
        qdrant_client=qdrant_client,
        embedder=embedder,
    )
    await storage.store("persona-1", "favorite_food", "любит пиццу")
    await storage.store("persona-1", "city", "живет в Челябинске")
    await storage.store("persona-2", "pet", "любит пиццу")

    response = await memory_test_client.delete("/api/memory/persona/persona-1")

    assert response.status_code == 200
    assert response.json() == {"session_id": "", "deleted_count": 2}

    persona_one_response = await memory_test_client.get("/api/memory/persona/persona-1")
    persona_two_response = await memory_test_client.get("/api/memory/persona/persona-2")
    assert persona_one_response.json()["memories"] == []
    assert len(persona_two_response.json()["memories"]) == 1


@pytest.mark.asyncio
async def test_memory_api_falls_back_to_session_identifier_for_unknown_session(
    memory_test_context: tuple[AsyncClient, FakeDatabase, FakeQdrantClient, FakeEmbedder],
) -> None:
    memory_test_client, _, qdrant_client, embedder = memory_test_context
    storage = MemoryStorage(
        qdrant_client=qdrant_client,
        embedder=embedder,
    )
    await storage.store("orphan-session", "favorite_food", "любит пиццу")

    response = await memory_test_client.get("/api/memory/orphan-session")

    assert response.status_code == 200
    assert response.json()["persona_id"] == "orphan-session"
    assert len(response.json()["memories"]) == 1