36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
from datetime import datetime, timezone
|
|
from typing import Set
|
|
|
|
from fastapi.testclient import TestClient
|
|
from great_ai import GreatAI
|
|
from great_ai.context import get_context
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class DatedInput(BaseModel):
|
|
when: datetime
|
|
tags: Set[int]
|
|
|
|
|
|
def test_tinydb_persists_json_compatible_model_values(isolated_context: None) -> None:
|
|
@GreatAI.create
|
|
def predict(data: DatedInput) -> int:
|
|
return len(data.tags)
|
|
|
|
with TestClient(predict.app) as client:
|
|
response = client.post(
|
|
"/predict", json={"data": {"when": "2026-01-01T00:00:00Z", "tags": [1, 2]}}
|
|
)
|
|
assert response.status_code == 200
|
|
trace_id = response.json()["trace_id"]
|
|
database = get_context().tracing_database
|
|
trace = database.get(trace_id)
|
|
assert trace is not None
|
|
assert trace.output == 2
|
|
trace.feedback = {"when": datetime(2026, 1, 1, tzinfo=timezone.utc)}
|
|
database.update(trace_id, trace)
|
|
updated = database.get(trace_id)
|
|
assert updated is not None
|
|
assert updated.feedback == {"when": "2026-01-01T00:00:00Z"}
|
|
database.save_batch([trace.model_copy(update={"trace_id": "batch"})])
|
|
assert database.get("batch") is not None
|