Kartik
117 lines
3.7 KiB
Python
117 lines
3.7 KiB
Python
"""Tests for _search.py — shared mem0 search API helper."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
def test_search_memories_returns_results():
|
|
from _search import search_memories
|
|
|
|
fake_results = [
|
|
{"id": "abc123", "memory": "Use Postgres for auth", "metadata": {"type": "decision"}},
|
|
{"id": "def456", "memory": "Never use floats for money", "metadata": {"type": "anti_pattern"}},
|
|
]
|
|
|
|
def mock_urlopen(req, timeout=None):
|
|
resp = MagicMock()
|
|
resp.read.return_value = json.dumps({"results": fake_results}).encode()
|
|
resp.__enter__ = lambda s: s
|
|
resp.__exit__ = MagicMock(return_value=False)
|
|
return resp
|
|
|
|
with patch("urllib.request.urlopen", side_effect=mock_urlopen):
|
|
results = search_memories("test-key", "user1", "proj1", "auth decisions")
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["id"] == "abc123"
|
|
|
|
|
|
def test_search_memories_with_metadata_type():
|
|
from _search import search_memories
|
|
|
|
captured_body = {}
|
|
|
|
def mock_urlopen(req, timeout=None):
|
|
captured_body.update(json.loads(req.data.decode()))
|
|
resp = MagicMock()
|
|
resp.read.return_value = json.dumps({"results": []}).encode()
|
|
resp.__enter__ = lambda s: s
|
|
resp.__exit__ = MagicMock(return_value=False)
|
|
return resp
|
|
|
|
with patch("urllib.request.urlopen", side_effect=mock_urlopen):
|
|
search_memories("key", "user", "proj", "query", metadata_type="decision")
|
|
|
|
filters = captured_body["filters"]
|
|
assert {"metadata": {"type": "decision"}} in filters["AND"]
|
|
|
|
|
|
def test_search_memories_handles_api_error():
|
|
from _search import search_memories
|
|
|
|
with patch("urllib.request.urlopen", side_effect=Exception("timeout")):
|
|
results = search_memories("key", "user", "proj", "query")
|
|
|
|
assert results == []
|
|
|
|
|
|
def test_search_memories_handles_list_response():
|
|
from _search import search_memories
|
|
|
|
fake_results = [{"id": "abc", "memory": "test"}]
|
|
|
|
def mock_urlopen(req, timeout=None):
|
|
resp = MagicMock()
|
|
resp.read.return_value = json.dumps(fake_results).encode()
|
|
resp.__enter__ = lambda s: s
|
|
resp.__exit__ = MagicMock(return_value=False)
|
|
return resp
|
|
|
|
with patch("urllib.request.urlopen", side_effect=mock_urlopen):
|
|
results = search_memories("key", "user", "proj", "query")
|
|
|
|
assert len(results) == 1
|
|
|
|
|
|
def test_search_memories_respects_top_k():
|
|
from _search import search_memories
|
|
|
|
captured_body = {}
|
|
|
|
def mock_urlopen(req, timeout=None):
|
|
captured_body.update(json.loads(req.data.decode()))
|
|
resp = MagicMock()
|
|
resp.read.return_value = json.dumps({"results": []}).encode()
|
|
resp.__enter__ = lambda s: s
|
|
resp.__exit__ = MagicMock(return_value=False)
|
|
return resp
|
|
|
|
with patch("urllib.request.urlopen", side_effect=mock_urlopen):
|
|
search_memories("key", "user", "proj", "query", top_k=5)
|
|
|
|
assert captured_body["top_k"] == 5
|
|
|
|
|
|
def test_search_memories_no_api_key_returns_empty():
|
|
from _search import search_memories
|
|
|
|
results = search_memories("", "user", "proj", "query")
|
|
assert results == []
|
|
|
|
|
|
def test_format_results_for_context():
|
|
from _search import format_results_for_context
|
|
|
|
memories = [
|
|
{"id": "abc12345-long-id", "memory": "Use Postgres for auth", "metadata": {"type": "decision"}},
|
|
{"id": "def67890-long-id", "memory": "JWT tokens expire in 1h", "metadata": {"type": "convention"}},
|
|
]
|
|
|
|
output = format_results_for_context(memories, heading="Relevant memories")
|
|
assert "Relevant memories" in output
|
|
assert "[decision]" in output
|
|
assert "Use Postgres for auth" in output
|
|
assert "abc12345" in output
|