283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""Unit tests for SQLAlchemy ORM models in agentkit.memory.models.
|
|
|
|
Covers the two public ORM model classes — EpisodeModel and ExperienceModel:
|
|
table mapping, column types/constraints, Python-side default callables,
|
|
explicit-value construction, and edge cases (empty/long/unicode values).
|
|
|
|
Note: memory/models.py contains SQLAlchemy ORM models (not Pydantic), so
|
|
Pydantic-specific scenarios (ValidationError / model_dump / model_validate)
|
|
are replaced by ORM equivalents (column introspection + default callable
|
|
invocation + attribute assignment).
|
|
|
|
DB-dependent helpers (create_episodic_session_factory,
|
|
create_experience_session_factory, ensure_episodic_table) are NOT tested
|
|
here — they require a live PostgreSQL + asyncpg connection.
|
|
"""
|
|
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import DateTime, Float, String, Text
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
|
|
from agentkit.memory.models import EpisodeModel, ExperienceModel
|
|
|
|
|
|
def _assert_uuid_default(id_default) -> None:
|
|
"""SQLAlchemy wraps Python-side default callables to receive an execution
|
|
context; pass None for callables that don't inspect it."""
|
|
val = id_default(None)
|
|
assert isinstance(val, str)
|
|
uuid.UUID(val) # raises ValueError if not a valid uuid4 string
|
|
|
|
|
|
def _assert_utc_default(ts_default) -> None:
|
|
ts = ts_default(None)
|
|
assert isinstance(ts, datetime)
|
|
assert ts.tzinfo is not None
|
|
assert ts.utcoffset() == timezone.utc.utcoffset(None)
|
|
|
|
|
|
def _assert_unique_ids(id_default, n: int = 100) -> None:
|
|
ids = {id_default(None) for _ in range(n)}
|
|
assert len(ids) == n
|
|
|
|
|
|
class TestEpisodeModel:
|
|
def test_tablename(self):
|
|
assert EpisodeModel.__tablename__ == "episodic_memories"
|
|
|
|
def test_column_types(self):
|
|
cols = EpisodeModel.__table__.c
|
|
assert isinstance(cols.id.type, String)
|
|
assert isinstance(cols.agent_name.type, String)
|
|
assert isinstance(cols.task_type.type, String)
|
|
assert isinstance(cols.input_summary.type, Text)
|
|
assert isinstance(cols.output_summary.type, Text)
|
|
assert isinstance(cols.outcome.type, String)
|
|
assert isinstance(cols.quality_score.type, Float)
|
|
assert isinstance(cols.reflection.type, Text)
|
|
assert isinstance(cols.embedding.type, Text)
|
|
assert isinstance(cols.metadata.type, JSONB)
|
|
assert isinstance(cols.created_at.type, DateTime)
|
|
|
|
def test_primary_key(self):
|
|
assert EpisodeModel.__table__.c.id.primary_key is True
|
|
|
|
def test_indexed_columns(self):
|
|
cols = EpisodeModel.__table__.c
|
|
assert cols.agent_name.index is True
|
|
assert cols.task_type.index is True
|
|
assert cols.created_at.index is True
|
|
|
|
def test_nullable_columns(self):
|
|
cols = EpisodeModel.__table__.c
|
|
assert cols.embedding.nullable is True
|
|
assert cols.metadata.nullable is True
|
|
|
|
def test_metadata_attribute_maps_to_metadata_column(self):
|
|
# Python attribute is `metadata_` (avoids clash with Base.metadata);
|
|
# the underlying DB column name is "metadata".
|
|
assert hasattr(EpisodeModel, "metadata_")
|
|
assert EpisodeModel.__table__.c.metadata.name == "metadata"
|
|
|
|
def test_scalar_defaults(self):
|
|
cols = EpisodeModel.__table__.c
|
|
assert cols.input_summary.default.arg == ""
|
|
assert cols.output_summary.default.arg == ""
|
|
assert cols.reflection.default.arg == ""
|
|
assert cols.outcome.default.arg == "success"
|
|
assert cols.quality_score.default.arg == 0.5
|
|
|
|
def test_callable_default_id_is_uuid_string(self):
|
|
_assert_uuid_default(EpisodeModel.__table__.c.id.default.arg)
|
|
|
|
def test_callable_default_created_at_is_utc(self):
|
|
_assert_utc_default(EpisodeModel.__table__.c.created_at.default.arg)
|
|
|
|
def test_default_id_is_unique_across_calls(self):
|
|
_assert_unique_ids(EpisodeModel.__table__.c.id.default.arg)
|
|
|
|
def test_construct_with_explicit_values(self):
|
|
ts = datetime(2025, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
|
|
m = EpisodeModel(
|
|
id="ep-1",
|
|
agent_name="agent-a",
|
|
task_type="analysis",
|
|
input_summary="in",
|
|
output_summary="out",
|
|
outcome="failure",
|
|
quality_score=0.9,
|
|
reflection="ref",
|
|
embedding="[0.1, 0.2]",
|
|
metadata_={"k": "v"},
|
|
created_at=ts,
|
|
)
|
|
assert m.id == "ep-1"
|
|
assert m.agent_name == "agent-a"
|
|
assert m.task_type == "analysis"
|
|
assert m.input_summary == "in"
|
|
assert m.output_summary == "out"
|
|
assert m.outcome == "failure"
|
|
assert m.quality_score == 0.9
|
|
assert m.reflection == "ref"
|
|
assert m.embedding == "[0.1, 0.2]"
|
|
assert m.metadata_ == {"k": "v"}
|
|
assert m.created_at == ts
|
|
|
|
def test_construct_minimal_yields_none_attrs(self):
|
|
# ORM Python-side defaults fire at flush (INSERT) time, not instantiation,
|
|
# so unset attributes read back as None.
|
|
m = EpisodeModel()
|
|
assert m.id is None
|
|
assert m.agent_name is None
|
|
assert m.created_at is None
|
|
assert m.metadata_ is None
|
|
|
|
def test_edge_empty_embedding_and_text(self):
|
|
m = EpisodeModel(embedding="", reflection="", input_summary="")
|
|
assert m.embedding == ""
|
|
assert m.reflection == ""
|
|
assert m.input_summary == ""
|
|
|
|
def test_edge_empty_vector_json(self):
|
|
m = EpisodeModel(embedding="[]")
|
|
assert m.embedding == "[]"
|
|
|
|
def test_edge_long_text(self):
|
|
long_text = "x" * 10_000
|
|
m = EpisodeModel(input_summary=long_text, output_summary=long_text, reflection=long_text)
|
|
assert len(m.input_summary) == 10_000
|
|
assert len(m.output_summary) == 10_000
|
|
assert len(m.reflection) == 10_000
|
|
|
|
def test_edge_unicode_and_special_characters(self):
|
|
m = EpisodeModel(
|
|
input_summary="中文 + emoji 🚀 + quote \" ' \n\t",
|
|
metadata_={"unicode": "你好", "null": None, "nested": {"a": [1, 2]}},
|
|
)
|
|
assert "🚀" in m.input_summary
|
|
assert "中文" in m.input_summary
|
|
assert m.metadata_["unicode"] == "你好"
|
|
assert m.metadata_["null"] is None
|
|
assert m.metadata_["nested"] == {"a": [1, 2]}
|
|
|
|
def test_edge_empty_metadata_dict(self):
|
|
m = EpisodeModel(metadata_={})
|
|
assert m.metadata_ == {}
|
|
|
|
|
|
class TestExperienceModel:
|
|
def test_tablename(self):
|
|
assert ExperienceModel.__tablename__ == "task_experiences"
|
|
|
|
def test_column_types(self):
|
|
cols = ExperienceModel.__table__.c
|
|
assert isinstance(cols.id.type, String)
|
|
assert isinstance(cols.task_type.type, String)
|
|
assert isinstance(cols.goal.type, Text)
|
|
assert isinstance(cols.steps_summary.type, Text)
|
|
assert isinstance(cols.outcome.type, String)
|
|
assert isinstance(cols.duration_seconds.type, Float)
|
|
assert isinstance(cols.success_rate.type, Float)
|
|
assert isinstance(cols.failure_reasons.type, JSONB)
|
|
assert isinstance(cols.optimization_tips.type, JSONB)
|
|
assert isinstance(cols.embedding.type, Text)
|
|
assert isinstance(cols.created_at.type, DateTime)
|
|
|
|
def test_primary_key_and_indexes(self):
|
|
cols = ExperienceModel.__table__.c
|
|
assert cols.id.primary_key is True
|
|
assert cols.task_type.index is True
|
|
assert cols.created_at.index is True
|
|
|
|
def test_nullable_embedding(self):
|
|
assert ExperienceModel.__table__.c.embedding.nullable is True
|
|
|
|
def test_scalar_defaults(self):
|
|
cols = ExperienceModel.__table__.c
|
|
assert cols.goal.default.arg == ""
|
|
assert cols.steps_summary.default.arg == ""
|
|
assert cols.outcome.default.arg == "success"
|
|
assert cols.duration_seconds.default.arg == 0.0
|
|
assert cols.success_rate.default.arg == 1.0
|
|
|
|
def test_callable_list_defaults(self):
|
|
cols = ExperienceModel.__table__.c
|
|
# default=list (the type itself, callable → empty list when invoked)
|
|
assert cols.failure_reasons.default.is_callable
|
|
assert cols.optimization_tips.default.is_callable
|
|
assert cols.failure_reasons.default.arg(None) == []
|
|
assert cols.optimization_tips.default.arg(None) == []
|
|
|
|
def test_callable_default_id_is_uuid_string(self):
|
|
_assert_uuid_default(ExperienceModel.__table__.c.id.default.arg)
|
|
|
|
def test_callable_default_created_at_is_utc(self):
|
|
_assert_utc_default(ExperienceModel.__table__.c.created_at.default.arg)
|
|
|
|
def test_default_id_is_unique_across_calls(self):
|
|
_assert_unique_ids(ExperienceModel.__table__.c.id.default.arg)
|
|
|
|
def test_construct_with_explicit_values(self):
|
|
ts = datetime(2025, 6, 7, 8, 9, 10, tzinfo=timezone.utc)
|
|
m = ExperienceModel(
|
|
id="exp-1",
|
|
task_type="build",
|
|
goal="ship feature",
|
|
steps_summary="step1; step2",
|
|
outcome="partial",
|
|
duration_seconds=42.5,
|
|
success_rate=0.66,
|
|
failure_reasons=["timeout", "oom"],
|
|
optimization_tips=["retry", "cache"],
|
|
embedding="[0.3, 0.4]",
|
|
created_at=ts,
|
|
)
|
|
assert m.id == "exp-1"
|
|
assert m.task_type == "build"
|
|
assert m.goal == "ship feature"
|
|
assert m.steps_summary == "step1; step2"
|
|
assert m.outcome == "partial"
|
|
assert m.duration_seconds == 42.5
|
|
assert m.success_rate == 0.66
|
|
assert m.failure_reasons == ["timeout", "oom"]
|
|
assert m.optimization_tips == ["retry", "cache"]
|
|
assert m.embedding == "[0.3, 0.4]"
|
|
assert m.created_at == ts
|
|
|
|
def test_construct_minimal_yields_none_attrs(self):
|
|
m = ExperienceModel()
|
|
assert m.id is None
|
|
assert m.task_type is None
|
|
assert m.created_at is None
|
|
assert m.failure_reasons is None
|
|
assert m.optimization_tips is None
|
|
|
|
def test_edge_empty_lists_and_embedding(self):
|
|
m = ExperienceModel(
|
|
failure_reasons=[],
|
|
optimization_tips=[],
|
|
embedding="",
|
|
)
|
|
assert m.failure_reasons == []
|
|
assert m.optimization_tips == []
|
|
assert m.embedding == ""
|
|
|
|
def test_edge_long_text(self):
|
|
long_text = "y" * 10_000
|
|
m = ExperienceModel(goal=long_text, steps_summary=long_text)
|
|
assert len(m.goal) == 10_000
|
|
assert len(m.steps_summary) == 10_000
|
|
|
|
def test_edge_unicode_and_special_characters(self):
|
|
m = ExperienceModel(
|
|
goal="中文 🎯 ' \" \n\t",
|
|
failure_reasons=["错误", "🚨"],
|
|
optimization_tips=["优化"],
|
|
)
|
|
assert "🎯" in m.goal
|
|
assert "中文" in m.goal
|
|
assert m.failure_reasons == ["错误", "🚨"]
|
|
assert m.optimization_tips == ["优化"]
|