7ef116e2f9
Define User, Project, Experiment, Run, StageResult, Score, ResponseCache, and WebhookConfig with UUID primary keys, JSON columns, enum types (ExperimentStatus, RunStatus), full relationship cascades, and indexes. Uses sqlalchemy.JSON (not JSONB) for SQLite compatibility in single-container mode. 16 tests added covering table creation, CRUD, uniqueness constraints, default values, and cascade deletes — all passing.
359 lines
10 KiB
Python
359 lines
10 KiB
Python
"""Tests for SQLAlchemy ORM models."""
|
|
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import create_engine, inspect
|
|
from sqlalchemy.orm import Session
|
|
|
|
from models import (
|
|
Base,
|
|
Experiment,
|
|
ExperimentStatus,
|
|
Project,
|
|
ResponseCache,
|
|
Run,
|
|
RunStatus,
|
|
Score,
|
|
StageResult,
|
|
User,
|
|
WebhookConfig,
|
|
)
|
|
|
|
|
|
def _engine():
|
|
engine = create_engine("sqlite:///:memory:")
|
|
Base.metadata.create_all(engine)
|
|
return engine
|
|
|
|
|
|
def _session(engine):
|
|
return Session(engine)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Table existence
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_all_tables_created():
|
|
engine = _engine()
|
|
table_names = inspect(engine).get_table_names()
|
|
expected = {
|
|
"users",
|
|
"projects",
|
|
"experiments",
|
|
"runs",
|
|
"stage_results",
|
|
"scores",
|
|
"response_cache",
|
|
"webhook_configs",
|
|
}
|
|
assert expected.issubset(set(table_names))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# User
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_user_creation():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="admin", password_hash="hashed", is_admin=True)
|
|
session.add(user)
|
|
session.commit()
|
|
|
|
assert isinstance(user.id, uuid.UUID)
|
|
assert user.username == "admin"
|
|
assert user.is_admin is True
|
|
assert isinstance(user.created_at, datetime)
|
|
|
|
|
|
def test_user_username_unique():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
session.add(User(username="dup", password_hash="h1"))
|
|
session.commit()
|
|
session.add(User(username="dup", password_hash="h2"))
|
|
try:
|
|
session.commit()
|
|
assert False, "Should have raised IntegrityError"
|
|
except Exception:
|
|
session.rollback()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Project
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_project_with_owner():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="owner", password_hash="h")
|
|
project = Project(name="Test Project", description="A test", owner=user)
|
|
session.add(project)
|
|
session.commit()
|
|
|
|
assert project.owner_id == user.id
|
|
assert project.name == "Test Project"
|
|
assert isinstance(project.updated_at, datetime)
|
|
|
|
|
|
def test_project_cascade_delete_from_user():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="owner", password_hash="h")
|
|
project = Project(name="P1", owner=user)
|
|
session.add(project)
|
|
session.commit()
|
|
project_id = project.id
|
|
|
|
session.delete(user)
|
|
session.commit()
|
|
|
|
assert session.get(Project, project_id) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Experiment
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_experiment_defaults():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(
|
|
project=project,
|
|
name="Exp1",
|
|
sample_data={"inputs": ["hello"]},
|
|
pipeline_stages=[{"prompt": "test"}],
|
|
scoring_config={"scorers": ["keyword"]},
|
|
parameter_space={"temperature": [0.1, 0.5]},
|
|
)
|
|
session.add(exp)
|
|
session.commit()
|
|
|
|
assert exp.status == ExperimentStatus.draft
|
|
assert exp.sample_data == {"inputs": ["hello"]}
|
|
assert isinstance(exp.created_at, datetime)
|
|
|
|
|
|
def test_experiment_cascade_delete_from_project():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
session.add(exp)
|
|
session.commit()
|
|
exp_id = exp.id
|
|
|
|
session.delete(project)
|
|
session.commit()
|
|
|
|
assert session.get(Experiment, exp_id) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Run
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_creation():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
run = Run(
|
|
experiment=exp,
|
|
config_hash="a" * 64,
|
|
config={"model": "gpt-4", "temperature": 0.5},
|
|
status=RunStatus.completed,
|
|
duration_ms=1200,
|
|
tokens_in=100,
|
|
tokens_out=50,
|
|
)
|
|
session.add(run)
|
|
session.commit()
|
|
|
|
assert run.status == RunStatus.completed
|
|
assert run.config["model"] == "gpt-4"
|
|
|
|
|
|
def test_run_default_status():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
run = Run(experiment=exp, config_hash="b" * 64, config={})
|
|
session.add(run)
|
|
session.commit()
|
|
|
|
assert run.status == RunStatus.pending
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# StageResult
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_stage_result():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
run = Run(experiment=exp, config_hash="c" * 64, config={})
|
|
sr = StageResult(
|
|
run=run,
|
|
stage_index=0,
|
|
prompt_sent="Hello",
|
|
response_raw="World",
|
|
model_used="gpt-4",
|
|
parameters={"temperature": 0.5},
|
|
tokens_in=10,
|
|
tokens_out=5,
|
|
latency_ms=200,
|
|
)
|
|
session.add(sr)
|
|
session.commit()
|
|
|
|
assert sr.stage_index == 0
|
|
assert sr.model_used == "gpt-4"
|
|
assert len(run.stage_results) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Score
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_score():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
run = Run(experiment=exp, config_hash="d" * 64, config={})
|
|
score = Score(
|
|
run=run,
|
|
scorer_name="embedding_similarity",
|
|
value=0.87,
|
|
scorer_metadata={"reference_id": "ref1"},
|
|
)
|
|
session.add(score)
|
|
session.commit()
|
|
|
|
assert score.value == 0.87
|
|
assert score.scorer_name == "embedding_similarity"
|
|
assert len(run.scores) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ResponseCache
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_response_cache():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
cache = ResponseCache(
|
|
config_hash="e" * 64,
|
|
response="cached response",
|
|
model="gpt-4",
|
|
tokens_in=50,
|
|
tokens_out=25,
|
|
latency_ms=300,
|
|
)
|
|
session.add(cache)
|
|
session.commit()
|
|
|
|
fetched = session.get(ResponseCache, "e" * 64)
|
|
assert fetched is not None
|
|
assert fetched.response == "cached response"
|
|
|
|
|
|
def test_response_cache_pk_is_config_hash():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
session.add(
|
|
ResponseCache(config_hash="f" * 64, response="r1", model="m1")
|
|
)
|
|
session.commit()
|
|
session.add(
|
|
ResponseCache(config_hash="f" * 64, response="r2", model="m2")
|
|
)
|
|
try:
|
|
session.commit()
|
|
assert False, "Should have raised IntegrityError"
|
|
except Exception:
|
|
session.rollback()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# WebhookConfig
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_webhook_config():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
wh = WebhookConfig(
|
|
event_type="experiment.completed",
|
|
url="https://example.com/hook",
|
|
headers={"Authorization": "Bearer token"},
|
|
is_active=True,
|
|
)
|
|
session.add(wh)
|
|
session.commit()
|
|
|
|
assert isinstance(wh.id, uuid.UUID)
|
|
assert wh.event_type == "experiment.completed"
|
|
assert wh.is_active is True
|
|
|
|
|
|
def test_webhook_config_default_active():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
wh = WebhookConfig(
|
|
event_type="run.failed",
|
|
url="https://example.com/hook",
|
|
)
|
|
session.add(wh)
|
|
session.commit()
|
|
|
|
assert wh.is_active is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Relationship cascades: Run → StageResult + Score
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_cascade_deletes_children():
|
|
engine = _engine()
|
|
with _session(engine) as session:
|
|
user = User(username="u", password_hash="h")
|
|
project = Project(name="P", owner=user)
|
|
exp = Experiment(project=project, name="E")
|
|
run = Run(experiment=exp, config_hash="g" * 64, config={})
|
|
sr = StageResult(
|
|
run=run, stage_index=0, prompt_sent="p",
|
|
response_raw="r", model_used="m",
|
|
)
|
|
score = Score(run=run, scorer_name="test", value=0.5)
|
|
session.add_all([run, sr, score])
|
|
session.commit()
|
|
|
|
sr_id, score_id = sr.id, score.id
|
|
session.delete(run)
|
|
session.commit()
|
|
|
|
assert session.get(StageResult, sr_id) is None
|
|
assert session.get(Score, score_id) is None
|