3ce89ea560
Co-authored-by: alvinfei <alvinfei@tencent.com> Co-authored-by: Cursor <cursoragent@cursor.com>
137 lines
3.9 KiB
Python
137 lines
3.9 KiB
Python
import asyncio
|
|
import threading
|
|
|
|
import pytest
|
|
|
|
from hy_memory.config import MemoryConfig
|
|
from hy_memory.data.cache_sqlite import SqliteCache
|
|
from hy_memory.data.history_store import HistoryStore
|
|
|
|
|
|
class _BlockingCursor:
|
|
def fetchone(self):
|
|
return (1,)
|
|
|
|
def fetchall(self):
|
|
return []
|
|
|
|
|
|
class _BlockingConnection:
|
|
"""Hold execute() open so a concurrent close() can be observed."""
|
|
|
|
def __init__(self):
|
|
self.entered = threading.Event()
|
|
self.release = threading.Event()
|
|
self.closed = threading.Event()
|
|
|
|
def execute(self, *_args, **_kwargs):
|
|
self.entered.set()
|
|
if not self.release.wait(timeout=2):
|
|
raise TimeoutError("test did not release the blocking query")
|
|
if self.closed.is_set():
|
|
raise AssertionError("connection was closed while execute() was active")
|
|
return _BlockingCursor()
|
|
|
|
def close(self):
|
|
self.closed.set()
|
|
|
|
|
|
def _config(tmp_path) -> MemoryConfig:
|
|
config = MemoryConfig()
|
|
config.history.db_path = str(tmp_path / "history.db")
|
|
config.cache.db_path = str(tmp_path / "cache.db")
|
|
return config
|
|
|
|
|
|
async def _wait_for(event: threading.Event) -> None:
|
|
assert await asyncio.wait_for(asyncio.to_thread(event.wait, 1), timeout=2)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_close_waits_for_active_query(tmp_path):
|
|
store = HistoryStore(_config(tmp_path))
|
|
connection = _BlockingConnection()
|
|
store._conn = connection
|
|
|
|
query = asyncio.create_task(store.count())
|
|
await _wait_for(connection.entered)
|
|
closing = asyncio.create_task(store.close())
|
|
await asyncio.sleep(0.05)
|
|
|
|
assert not connection.closed.is_set()
|
|
connection.release.set()
|
|
assert await query == 1
|
|
await closing
|
|
assert connection.closed.is_set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_close_waits_for_active_query(tmp_path):
|
|
store = SqliteCache(_config(tmp_path))
|
|
connection = _BlockingConnection()
|
|
store._conn = connection
|
|
|
|
query = asyncio.create_task(store.load_metrics_range("0", "9"))
|
|
await _wait_for(connection.entered)
|
|
closing = asyncio.create_task(store.close())
|
|
await asyncio.sleep(0.05)
|
|
|
|
assert not connection.closed.is_set()
|
|
connection.release.set()
|
|
assert await query == []
|
|
await closing
|
|
assert connection.closed.is_set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_history_concurrent_reads_and_writes_are_serialized(tmp_path):
|
|
store = HistoryStore(_config(tmp_path))
|
|
await store.initialize()
|
|
try:
|
|
inserted = await asyncio.gather(
|
|
*(
|
|
store.record_add(
|
|
memory_id=f"memory-{index}",
|
|
content="content",
|
|
layer="l2_fact",
|
|
isolation_key="user",
|
|
)
|
|
for index in range(200)
|
|
)
|
|
)
|
|
assert all(row_id > 0 for row_id in inserted)
|
|
assert len(set(inserted)) == 200
|
|
|
|
counts = await asyncio.gather(*(store.count() for _ in range(100)))
|
|
assert counts == [200] * 100
|
|
finally:
|
|
await store.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_concurrent_pipeline_log_reads_and_writes_are_serialized(tmp_path):
|
|
store = SqliteCache(_config(tmp_path))
|
|
await store.initialize()
|
|
try:
|
|
written = await asyncio.gather(
|
|
*(
|
|
store.store_pipeline_log(
|
|
request_id="request",
|
|
user_id="user",
|
|
agent_id="agent",
|
|
step="EXTRACT",
|
|
prompt=f"prompt-{index}",
|
|
response="response",
|
|
)
|
|
for index in range(200)
|
|
)
|
|
)
|
|
assert written == [True] * 200
|
|
|
|
pages = await asyncio.gather(
|
|
*(store.get_pipeline_logs(request_id="request", limit=200) for _ in range(50))
|
|
)
|
|
assert all(len(page) == 200 for page in pages)
|
|
finally:
|
|
await store.close()
|