wehub-resource-sync
411 lines
14 KiB
Python
411 lines
14 KiB
Python
"""Tests for AsyncSQLiteSession functionality."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
from collections.abc import Sequence
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("aiosqlite") # Skip tests if aiosqlite is not installed
|
|
|
|
from agents import Agent, Runner, TResponseInputItem
|
|
from agents.extensions.memory import AsyncSQLiteSession
|
|
from agents.memory import SessionSettings
|
|
from tests.fake_model import FakeModel
|
|
from tests.test_responses import get_text_message
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
@pytest.fixture
|
|
def agent() -> Agent:
|
|
"""Fixture for a basic agent with a fake model."""
|
|
return Agent(name="test", model=FakeModel())
|
|
|
|
|
|
def _item_ids(items: Sequence[TResponseInputItem]) -> list[str]:
|
|
result: list[str] = []
|
|
for item in items:
|
|
item_dict = cast(dict[str, Any], item)
|
|
result.append(cast(str, item_dict["id"]))
|
|
return result
|
|
|
|
|
|
async def test_async_sqlite_session_basic_flow():
|
|
"""Test AsyncSQLiteSession add/get/clear behavior."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_basic.db"
|
|
session = AsyncSQLiteSession("async_basic", db_path)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
|
|
await session.add_items(items)
|
|
retrieved = await session.get_items()
|
|
assert retrieved == items
|
|
|
|
await session.clear_session()
|
|
assert await session.get_items() == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_pop_item():
|
|
"""Test AsyncSQLiteSession pop_item behavior."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_pop.db"
|
|
session = AsyncSQLiteSession("async_pop", db_path)
|
|
|
|
assert await session.pop_item() is None
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "One"},
|
|
{"role": "assistant", "content": "Two"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
popped = await session.pop_item()
|
|
assert popped == items[-1]
|
|
assert await session.get_items() == items[:-1]
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_pop_item_skips_corrupt_most_recent():
|
|
"""pop_item skips corrupt newest rows and returns the next valid item."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_pop_corrupt.db"
|
|
session = AsyncSQLiteSession("async_pop_corrupt", db_path)
|
|
|
|
valid_item: TResponseInputItem = {"role": "user", "content": "valid"}
|
|
await session.add_items([valid_item])
|
|
|
|
conn = await session._get_connection()
|
|
await conn.execute(
|
|
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
|
|
(session.session_id, "not valid json {{{"),
|
|
)
|
|
await conn.commit()
|
|
|
|
assert await session.pop_item() == valid_item
|
|
assert await session.get_items() == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_pop_item_returns_none_after_dropping_only_corrupt_rows():
|
|
"""pop_item removes corrupt rows and returns None when no valid items remain."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_pop_only_corrupt.db"
|
|
session = AsyncSQLiteSession("async_pop_only_corrupt", db_path)
|
|
|
|
conn = await session._get_connection()
|
|
await conn.execute(
|
|
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
|
|
(session.session_id, "not valid json {{{"),
|
|
)
|
|
await conn.commit()
|
|
|
|
assert await session.pop_item() is None
|
|
assert await session.get_items() == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_get_items_limit():
|
|
"""Test AsyncSQLiteSession get_items limit handling."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_limit.db"
|
|
session = AsyncSQLiteSession("async_limit", db_path)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "Message 1"},
|
|
{"role": "assistant", "content": "Response 1"},
|
|
{"role": "user", "content": "Message 2"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
latest = await session.get_items(limit=2)
|
|
assert latest == items[-2:]
|
|
|
|
none = await session.get_items(limit=0)
|
|
assert none == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_session_settings_default():
|
|
"""Test that session_settings defaults to empty SessionSettings."""
|
|
session = AsyncSQLiteSession("async_default_settings")
|
|
|
|
assert isinstance(session.session_settings, SessionSettings)
|
|
assert session.session_settings.limit is None
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_session_settings_constructor():
|
|
"""Test passing session_settings via constructor."""
|
|
session = AsyncSQLiteSession(
|
|
"async_constructor_settings",
|
|
session_settings=SessionSettings(limit=5),
|
|
)
|
|
|
|
assert session.session_settings is not None
|
|
assert session.session_settings.limit == 5
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_get_items_uses_session_settings_limit():
|
|
"""Test that get_items uses session_settings.limit as default."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_settings_limit.db"
|
|
session = AsyncSQLiteSession(
|
|
"async_settings_limit",
|
|
db_path,
|
|
session_settings=SessionSettings(limit=3),
|
|
)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": f"Message {i}"} for i in range(5)
|
|
]
|
|
await session.add_items(items)
|
|
|
|
retrieved = await session.get_items()
|
|
assert retrieved == items[-3:]
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_explicit_limit_overrides_session_settings():
|
|
"""Test that explicit limit parameter overrides session_settings."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_settings_override.db"
|
|
session = AsyncSQLiteSession(
|
|
"async_settings_override",
|
|
db_path,
|
|
session_settings=SessionSettings(limit=5),
|
|
)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": f"Message {i}"} for i in range(10)
|
|
]
|
|
await session.add_items(items)
|
|
|
|
retrieved = await session.get_items(limit=2)
|
|
assert retrieved == items[-2:]
|
|
|
|
no_items = await session.get_items(limit=0)
|
|
assert no_items == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_unicode_content():
|
|
"""Test AsyncSQLiteSession stores unicode content."""
|
|
session = AsyncSQLiteSession("async_unicode")
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "こんにちは"},
|
|
{"role": "assistant", "content": "Привет"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
retrieved = await session.get_items()
|
|
assert retrieved == items
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_runner_integration(agent: Agent):
|
|
"""Test that AsyncSQLiteSession works correctly with the agent Runner."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_runner_integration.db"
|
|
session = AsyncSQLiteSession("runner_integration_test", db_path)
|
|
|
|
assert isinstance(agent.model, FakeModel)
|
|
|
|
agent.model.set_next_output([get_text_message("San Francisco")])
|
|
result1 = await Runner.run(
|
|
agent,
|
|
"What city is the Golden Gate Bridge in?",
|
|
session=session,
|
|
)
|
|
assert result1.final_output == "San Francisco"
|
|
|
|
agent.model.set_next_output([get_text_message("California")])
|
|
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
|
assert result2.final_output == "California"
|
|
|
|
last_input = agent.model.last_turn_args["input"]
|
|
assert isinstance(last_input, list)
|
|
assert len(last_input) > 1
|
|
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_session_isolation(agent: Agent):
|
|
"""Test that different session IDs result in isolated conversation histories."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_isolation.db"
|
|
session1 = AsyncSQLiteSession("session_1", db_path)
|
|
session2 = AsyncSQLiteSession("session_2", db_path)
|
|
|
|
assert isinstance(agent.model, FakeModel)
|
|
agent.model.set_next_output([get_text_message("I like cats.")])
|
|
await Runner.run(agent, "I like cats.", session=session1)
|
|
|
|
agent.model.set_next_output([get_text_message("I like dogs.")])
|
|
await Runner.run(agent, "I like dogs.", session=session2)
|
|
|
|
agent.model.set_next_output([get_text_message("You said you like cats.")])
|
|
result = await Runner.run(agent, "What animal did I say I like?", session=session1)
|
|
assert "cats" in result.final_output.lower()
|
|
assert "dogs" not in result.final_output.lower()
|
|
|
|
await session1.close()
|
|
await session2.close()
|
|
|
|
|
|
async def test_async_sqlite_session_add_empty_items_list():
|
|
"""Test that adding an empty list of items is a no-op."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_add_empty.db"
|
|
session = AsyncSQLiteSession("add_empty_test", db_path)
|
|
|
|
assert await session.get_items() == []
|
|
await session.add_items([])
|
|
assert await session.get_items() == []
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_pop_from_empty_session():
|
|
"""Test that pop_item returns None on an empty session."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_pop_empty.db"
|
|
session = AsyncSQLiteSession("empty_session", db_path)
|
|
|
|
popped = await session.pop_item()
|
|
assert popped is None
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_get_items_with_limit_more_than_available():
|
|
"""Test limit behavior when requesting more items than exist."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_limit_more.db"
|
|
session = AsyncSQLiteSession("limit_more_test", db_path)
|
|
|
|
items: list[TResponseInputItem] = [
|
|
{"role": "user", "content": "1"},
|
|
{"role": "assistant", "content": "2"},
|
|
{"role": "user", "content": "3"},
|
|
{"role": "assistant", "content": "4"},
|
|
]
|
|
await session.add_items(items)
|
|
|
|
retrieved = await session.get_items(limit=10)
|
|
assert retrieved == items
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_get_items_same_timestamp_consistent_order():
|
|
"""Test that items with identical timestamps keep insertion order."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_same_timestamp.db"
|
|
session = AsyncSQLiteSession("same_timestamp_test", db_path)
|
|
|
|
older_item = cast(
|
|
TResponseInputItem, {"id": "older_same_ts", "role": "user", "content": "old"}
|
|
)
|
|
reasoning_item = cast(TResponseInputItem, {"id": "rs_same_ts", "type": "reasoning"})
|
|
message_item = cast(
|
|
TResponseInputItem,
|
|
{"id": "msg_same_ts", "type": "message", "role": "assistant", "content": []},
|
|
)
|
|
|
|
await session.add_items([older_item])
|
|
await session.add_items([reasoning_item, message_item])
|
|
|
|
conn = await session._get_connection()
|
|
cursor = await conn.execute(
|
|
f"SELECT id, message_data FROM {session.messages_table} WHERE session_id = ?",
|
|
(session.session_id,),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
await cursor.close()
|
|
|
|
id_map: dict[str, int] = {
|
|
cast(str, json.loads(message_json)["id"]): cast(int, row_id)
|
|
for row_id, message_json in rows
|
|
}
|
|
|
|
shared = datetime(2025, 10, 15, 17, 26, 39, 132483)
|
|
shared_str = shared.strftime("%Y-%m-%d %H:%M:%S.%f")
|
|
await conn.execute(
|
|
f"""
|
|
UPDATE {session.messages_table}
|
|
SET created_at = ?
|
|
WHERE id IN (?, ?, ?)
|
|
""",
|
|
(
|
|
shared_str,
|
|
id_map["older_same_ts"],
|
|
id_map["rs_same_ts"],
|
|
id_map["msg_same_ts"],
|
|
),
|
|
)
|
|
await conn.commit()
|
|
|
|
retrieved = await session.get_items()
|
|
assert _item_ids(retrieved) == ["older_same_ts", "rs_same_ts", "msg_same_ts"]
|
|
|
|
latest_two = await session.get_items(limit=2)
|
|
assert _item_ids(latest_two) == ["rs_same_ts", "msg_same_ts"]
|
|
|
|
await session.close()
|
|
|
|
|
|
async def test_async_sqlite_session_pop_item_same_timestamp_returns_latest():
|
|
"""Test that pop_item returns the newest item when timestamps tie."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "async_same_timestamp_pop.db"
|
|
session = AsyncSQLiteSession("same_timestamp_pop_test", db_path)
|
|
|
|
reasoning_item = cast(TResponseInputItem, {"id": "rs_pop_same_ts", "type": "reasoning"})
|
|
message_item = cast(
|
|
TResponseInputItem,
|
|
{"id": "msg_pop_same_ts", "type": "message", "role": "assistant", "content": []},
|
|
)
|
|
|
|
await session.add_items([reasoning_item, message_item])
|
|
|
|
conn = await session._get_connection()
|
|
shared = datetime(2025, 10, 15, 17, 26, 39, 132483)
|
|
shared_str = shared.strftime("%Y-%m-%d %H:%M:%S.%f")
|
|
await conn.execute(
|
|
f"UPDATE {session.messages_table} SET created_at = ? WHERE session_id = ?",
|
|
(shared_str, session.session_id),
|
|
)
|
|
await conn.commit()
|
|
|
|
popped = await session.pop_item()
|
|
assert popped is not None
|
|
assert cast(dict[str, Any], popped)["id"] == "msg_pop_same_ts"
|
|
|
|
remaining = await session.get_items()
|
|
assert _item_ids(remaining) == ["rs_pop_same_ts"]
|
|
|
|
await session.close()
|