"""Unit tests for MemoryStore (mocked asyncpg)."""
import re
import pytest
from navi.config import settings
from navi.memory._facts import _key_tokens
from tests.conftest_factory import FakeConnection, FakeRecord, FakePool, make_store_with_pool
class TestKeyTokens:
"""`preferred_language`, `ui_language` and `language` are one key while `dark`
and `dark_ui` are two — hence set equality, not overlap.
"""
def test_prefixes_are_dropped(self):
assert _key_tokens("preferred_language") == _key_tokens("language")
assert _key_tokens("ui_language") == _key_tokens("language")
assert _key_tokens("response_language") == _key_tokens("language")
def test_word_order_of_meaningful_tokens_does_not_matter(self):
assert _key_tokens("primary_editor") == _key_tokens("editor_primary")
def test_a_trailing_noise_word_survives(self):
"""Noise is dropped from the front only. `ui_language` loses its `ui`;
`dark_ui` keeps it, which is what keeps that key apart from `dark`. The
cost is that `ui_dark` does lose it and does collapse into `dark`."""
assert _key_tokens("ui_language") == _key_tokens("language")
assert _key_tokens("dark_ui") == frozenset({"dark", "ui"})
assert _key_tokens("ui_dark") == frozenset({"dark"})
def test_overlap_without_equality_is_not_a_match(self):
assert _key_tokens("dark") != _key_tokens("dark_ui")
assert _key_tokens("primary_os") != _key_tokens("os")
def test_a_key_of_pure_noise_has_no_tokens(self):
assert _key_tokens("") == frozenset()
assert _key_tokens("ui_") == frozenset()
class TestUpsertFact:
async def test_calls_execute(self):
conn = FakeConnection()
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="profile", key="name", value="Eugene")
assert conn.calls[0][0] == "execute"
assert "memory_facts" in conn.calls[0][1]
async def test_without_embedding(self):
"""The no-embedding branch bound one placeholder more than it had columns
(13 vs 12), so every fact written while the embedding backend was down
died on PostgresSyntaxError and the extraction was lost — 47 embed
failures in three days on prod, most of them reaching this branch.
"""
conn = FakeConnection()
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
store._embedding_backend = None
await store.upsert_fact(category="profile", key="name", value="Eugene")
_, query, args = conn.calls[0]
columns = query.split("(", 1)[1].split(")", 1)[0].split(",")
highest_placeholder = max(int(n) for n in re.findall(r"\$(\d+)", query))
assert "embedding" not in query
assert highest_placeholder == len(columns) == len(args)
class TestExpirySurvivesAnExtractorWrite:
"""`SET expires_at = EXCLUDED.expires_at` let the extractor, which always passes
None, wipe an expiry the agent had set on the same key.
"""
async def test_plain_branch_coalesces(self):
conn = FakeConnection()
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="technical", key="ip", value="10.0.0.1")
assert "COALESCE(EXCLUDED.expires_at, memory_facts.expires_at)" in conn.calls[0][1]
async def test_embedding_branch_coalesces_too(self):
conn = FakeConnection()
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = True
store._embedding_backend = object()
async def _fake_embed(text: str):
return [0.1] * 768
store._generate_embedding = _fake_embed
await store.upsert_fact(category="technical", key="ip", value="10.0.0.1")
assert "COALESCE(EXCLUDED.expires_at, memory_facts.expires_at)" in conn.calls[0][1]
class TestMergeKeyAliases:
"""Key drift is what the extractor produces on its own: it writes the same fact
under a new spelling, and the exact-key upsert cannot see the old one. Prod had
11 such clusters / 24 rows out of 1017 facts.
"""
async def test_plain_upsert_never_looks_for_aliases(self):
conn = FakeConnection()
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="profile", key="name", value="Eugene")
assert [c[0] for c in conn.calls] == ["execute"]
async def test_marks_the_older_spelling(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="old-1", key="language", merged_into=None)])
conn.enqueue("UPDATE 1") # the mark, written before the row it points at
conn.enqueue("INSERT 0 1") # the new row itself
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="preferred_language",
value="ru", dedupe=True)
assert [c[0] for c in conn.calls] == ["fetch", "execute", "execute"]
_, query, args = conn.calls[1]
assert "SET merged_into" in query
inserted_id = conn.calls[2][2][0] # the id the INSERT carries
assert args == (["old-1"], inserted_id)
async def test_repoints_the_losers_of_an_earlier_merge(self):
"""Otherwise the loser of a loser keeps pointing at a row that is itself hidden."""
conn = FakeConnection()
conn.enqueue([FakeRecord(id="old-1", key="language", merged_into=None)])
conn.enqueue("UPDATE 2")
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="preferred_language",
value="ru", dedupe=True)
assert "merged_into = ANY($1::text[])" in conn.calls[1][1]
async def test_an_exact_key_stays_the_winner(self):
"""An exact key is the ON CONFLICT target, so the freshest value lands in it
and only the other spelling is hidden."""
conn = FakeConnection()
conn.enqueue([
FakeRecord(id="exact-1", key="language", merged_into=None),
FakeRecord(id="old-1", key="preferred_language", merged_into=None),
])
conn.enqueue("UPDATE 1")
conn.enqueue("UPDATE 1")
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="language",
value="ru", dedupe=True)
assert conn.calls[2][2] == (["old-1"], "exact-1")
async def test_an_exact_key_that_is_hidden_is_unhidden(self):
"""`language` was hidden behind `preferred_language`, and now the extractor
writes `language` again. The upsert lands on that hidden row through
ON CONFLICT — the id we generated is never inserted — so unless it is
un-hidden the write goes somewhere no read can see it, and the loser is
pointed at a row that does not exist.
"""
conn = FakeConnection()
conn.enqueue([
FakeRecord(id="exact-1", key="language", merged_into="live-1"),
FakeRecord(id="live-1", key="preferred_language", merged_into=None),
])
conn.enqueue("UPDATE 1") # un-hide the row the upsert will land on
conn.enqueue("UPDATE 1") # hide the spelling that was live
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="language",
value="ru", dedupe=True)
unhide, mark = conn.calls[1], conn.calls[2]
assert "SET merged_into = NULL" in unhide[1]
assert unhide[2] == (["exact-1"],)
assert mark[2] == (["live-1"], "exact-1")
async def test_a_hidden_row_is_not_hidden_twice(self):
"""The losers of an earlier merge are already behind something; only the
live spellings are the ones this write has to hide."""
conn = FakeConnection()
conn.enqueue([
FakeRecord(id="hidden-1", key="response_language", merged_into="live-1"),
FakeRecord(id="live-1", key="preferred_language", merged_into=None),
])
conn.enqueue("UPDATE 1")
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="language", value="ru",
dedupe=True)
_, _, args = conn.calls[1]
assert args[0] == ["live-1"] # hidden-1 is re-pointed by the second clause
async def test_a_different_key_is_left_alone(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="old-1", key="dark", merged_into=None)])
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="preferences", key="dark_ui", value="true",
dedupe=True)
assert [c[0] for c in conn.calls] == ["fetch", "execute"]
async def test_candidates_are_scoped_to_user_and_category(self):
conn = FakeConnection()
conn.enqueue([])
conn.enqueue("INSERT 0 1")
store = make_store_with_pool(conn)
await store.upsert_fact(category="technical", key="os", value="Arch",
user_id="u1", dedupe=True)
_, query, args = conn.calls[0]
assert "user_id IS NOT DISTINCT FROM $1" in query
assert args[:2] == ("u1", "technical")
class TestMergeFilterOnReads:
"""A merged loser that some read can still see is no better than a deleted row
nobody can undelete — the mark only means something while it hides the row.
"""
async def test_vector_search_hides_merged(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="", distance=0.1)])
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = True
store._embedding_backend = object()
async def _fake_embed(text: str):
return [0.1] * 768
store._generate_embedding = _fake_embed
await store.search_facts("name")
assert "merged_into IS NULL" in conn.calls[0][1]
async def test_ilike_fallback_hides_merged_and_expired(self):
"""This branch was also the one read path that served expired facts."""
conn = FakeConnection()
conn.enqueue(100) # fact_count, above the auto-dump threshold
conn.enqueue([])
conn.enqueue([])
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = False
await store.search_facts("eugene")
for call in (conn.calls[1], conn.calls[2]):
assert "merged_into IS NULL" in call[1]
assert "expires_at" in call[1]
async def test_get_all_facts_hides_merged_unless_asked(self):
conn = FakeConnection()
conn.enqueue([])
store = make_store_with_pool(conn)
await store.get_all_facts()
assert "merged_into IS NULL" in conn.calls[0][1]
conn2 = FakeConnection()
conn2.enqueue([])
store2 = make_store_with_pool(conn2)
await store2.get_all_facts(include_merged=True)
# The column is selected either way — the admin route needs to see the mark
# to undo a merge. Only the WHERE changes.
assert "merged_into" in conn2.calls[0][1]
assert "merged_into IS NULL" not in conn2.calls[0][1]
async def test_categories_hide_merged(self):
conn = FakeConnection()
conn.enqueue([])
store = make_store_with_pool(conn)
await store.get_categories()
assert "merged_into IS NULL" in conn.calls[0][1]
async def test_backfill_skips_merged(self):
conn = FakeConnection()
conn.enqueue([])
store = make_store_with_pool(conn)
await store.backfill_embeddings()
assert "merged_into IS NULL" in conn.calls[0][1]
class TestEmbeddingInputLength:
"""An input past the model's window came back as a 400 and took the whole
embedding with it — 41 reaction-session prompts did that on prod, each one
dropping recall back to ILIKE. Clip the head instead of losing the vector.
"""
class _RecordingBackend:
def __init__(self):
self.texts: list[str] = []
async def embed(self, texts, model=None):
self.texts = list(texts)
return [[0.1] * 768 for _ in texts]
def _store(self):
store = make_store_with_pool(FakeConnection())
store._pgvector_checked = True
store._pgvector_available = True
store._embedding_backend = self._RecordingBackend()
return store
async def test_single_text_is_clipped_before_the_backend_sees_it(self):
store = self._store()
over = settings.embedding_max_chars + 500
await store._generate_embedding("x" * over)
assert len(store._embedding_backend.texts[0]) == settings.embedding_max_chars
async def test_batch_clips_each_text_and_leaves_short_ones_alone(self):
store = self._store()
await store._generate_embeddings(["y" * (settings.embedding_max_chars * 2), "short"])
assert [len(t) for t in store._embedding_backend.texts] == [
settings.embedding_max_chars,
len("short"),
]
class TestSearchFacts:
async def test_vector_search_happy_path(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="", distance=0.1)])
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = True
store._embedding_backend = object() # any truthy object
# Mock _generate_embedding to avoid hitting the backend
async def _fake_embed(text: str):
return [0.1] * 768
store._generate_embedding = _fake_embed
results = await store.search_facts("name", limit=5)
assert len(results) == 1
assert results[0]["key"] == "name"
async def test_fallback_to_ilike_no_pgvector(self):
conn = FakeConnection()
conn.enqueue(0) # COUNT(*)
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="")])
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = False
results = await store.search_facts("eugene", limit=5)
assert len(results) == 1
async def test_fallback_auto_dump_below_threshold(self):
conn = FakeConnection()
conn.enqueue(5) # fact_count <= threshold
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="")])
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = False
results = await store.search_facts("anything", limit=5)
assert len(results) == 1
# Should have done get_all_facts instead of ILIKE
assert "ORDER BY category" in conn.calls[1][1]
async def test_fallback_no_terms(self):
conn = FakeConnection()
conn.enqueue([]) # get_all_facts returns empty
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = False
results = await store.search_facts("a", limit=5)
# single-char query normalizes to empty -> get_all_facts
assert len(results) == 0
class TestDeleteFact:
async def test_by_key(self):
conn = FakeConnection()
conn.enqueue("DELETE 0") # hidden aliases of the key
conn.enqueue("DELETE 1")
store = make_store_with_pool(conn)
count = await store.delete_fact("name")
assert count == 1
assert "DELETE FROM memory_facts" in conn.calls[1][1]
async def test_by_key_returns_zero(self):
conn = FakeConnection()
conn.enqueue("DELETE 0")
conn.enqueue("DELETE 0")
store = make_store_with_pool(conn)
count = await store.delete_fact("missing")
assert count == 0
async def test_by_key_and_category(self):
conn = FakeConnection()
conn.enqueue("DELETE 0")
conn.enqueue("DELETE 0 1")
store = make_store_with_pool(conn)
count = await store.delete_fact("name", category="profile")
assert "category" in conn.calls[1][1]
async def test_purges_the_aliases_of_the_deleted_key(self):
"""A merge loser points at its winner. Deleting the winner alone would leave
those rows hidden with nothing to point at and no way back to them.
"""
conn = FakeConnection()
conn.enqueue("DELETE 2")
conn.enqueue("DELETE 1")
store = make_store_with_pool(conn)
await store.delete_fact("language", category="preferences")
op, query, _args = conn.calls[0]
assert op == "execute"
assert "merged_into IN (SELECT id FROM memory_facts" in query
class TestGetAllFacts:
async def test_returns_records(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="")])
store = make_store_with_pool(conn)
results = await store.get_all_facts()
assert len(results) == 1
assert results[0]["key"] == "name"
async def test_with_limit(self):
conn = FakeConnection()
conn.enqueue([])
store = make_store_with_pool(conn)
await store.get_all_facts(limit=5)
assert "LIMIT $1" in conn.calls[0][1]
async def test_all_users(self):
conn = FakeConnection()
conn.enqueue([FakeRecord(id="1", category="profile", key="name", value="Eugene",
updated_at=None, source="conversation", confidence=90,
expires_at=None, source_context="")])
store = make_store_with_pool(conn)
results = await store.get_all_facts(all_users=True)
assert len(results) == 1
# Should not filter by user_id at all
assert "user_id" not in conn.calls[0][1]
class TestFactCount:
async def test_returns_count(self):
conn = FakeConnection()
conn.enqueue(42)
store = make_store_with_pool(conn)
assert await store.fact_count() == 42
async def test_all_users(self):
conn = FakeConnection()
conn.enqueue(100)
store = make_store_with_pool(conn)
assert await store.fact_count(all_users=True) == 100
# No user filter, but merged losers are still counted out of the live total.
assert "user_id" not in conn.calls[0][1]
assert "merged_into IS NULL" in conn.calls[0][1]
async def test_include_merged(self):
conn = FakeConnection()
conn.enqueue(100)
store = make_store_with_pool(conn)
assert await store.fact_count(include_merged=True) == 100
assert "merged_into" not in conn.calls[0][1]
class TestSummary:
async def test_get_summary(self):
conn = FakeConnection()
conn.enqueue("User likes Python.")
store = make_store_with_pool(conn)
assert await store.get_summary() == "User likes Python."
async def test_set_summary(self):
conn = FakeConnection()
conn.enqueue("OK")
store = make_store_with_pool(conn)
await store.set_summary("New summary")
assert "memory_summary" in conn.calls[0][1]
class TestSessionState:
async def test_mark_extracted(self):
conn = FakeConnection()
conn.enqueue("OK")
store = make_store_with_pool(conn)
await store.mark_session_extracted("sess-1")
assert "session_memory_state" in conn.calls[0][1]
async def test_get_extracted_at(self):
from datetime import datetime, timezone
now = datetime.now(timezone.utc)
conn = FakeConnection()
conn.enqueue(FakeRecord(extracted_at=now))
store = make_store_with_pool(conn)
result = await store.get_extracted_at("sess-1")
assert result == now.isoformat()
async def test_get_extracted_at_none(self):
conn = FakeConnection()
conn.enqueue(None)
store = make_store_with_pool(conn)
assert await store.get_extracted_at("sess-1") is None
class TestBackfillEmbeddings:
async def test_updates_rows(self):
conn = FakeConnection()
# First batch: 2 rows
conn.enqueue([
FakeRecord(id="1", value="hello"),
FakeRecord(id="2", value="world"),
])
# No more rows
conn.enqueue([])
# executemany response
conn.enqueue(None)
store = make_store_with_pool(conn)
store._pgvector_checked = True
store._pgvector_available = True
store._embedding_backend = object()
async def _fake_embeds(texts: list[str]):
return [[0.1] * 768 for _ in texts]
store._generate_embeddings = _fake_embeds
updated = await store.backfill_embeddings(batch_size=2)
assert updated == 2
assert conn.calls[-2][0] == "executemany"
assert "UPDATE memory_facts SET embedding" in conn.calls[-2][1]
