Tests

tests/test_cache_regressions.py

@@ -1,8 +1,10 @@
 import asyncio
-from typing import Dict, List, Tuple
+from typing import Any, Dict, List, Tuple
 
 import pytest
+from fastapi.testclient import TestClient
 from great_ai import GreatAI
+from great_ai.context import get_context
 
 
 def test_nested_tuple_inputs_can_be_cached(isolated_context: None) -> None:
@@ -73,3 +75,126 @@ def test_caller_mutation_does_not_change_cached_trace_values(
     assert predict(original).output == 2
     assert predict([[1]]).trace_id == first.trace_id
     assert first.logged_values["arg:values:value"] == [[1]]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.parametrize("do_not_persist", [False, True])
+async def test_returned_traces_cannot_mutate_cached_results(
+    isolated_context: None, is_async: bool, do_not_persist: bool
+) -> None:
+    calls = 0
+
+    def predict(values: List[int]) -> List[List[int]]:
+        nonlocal calls
+        calls += 1
+        return [values]
+
+    async def predict_async(values: List[int]) -> List[List[int]]:
+        return predict(values)
+
+    service: Any = GreatAI.create(predict_async if is_async else predict)
+    first = service([1], do_not_persist_traces=do_not_persist)
+    if is_async:
+        first = await first
+    trace_id = first.trace_id
+    first.output[0].append(99)
+    first.logged_values["arg:values:value"].append(88)
+    first.tags.append("changed")
+    first.trace_id = "changed"
+    first.feedback = "changed"
+
+    second = service([1], do_not_persist_traces=do_not_persist)
+    if is_async:
+        second = await second
+    assert second.trace_id == trace_id
+    assert second.output == [[1]]
+    assert second.logged_values["arg:values:value"] == [1]
+    assert "changed" not in second.tags
+    assert second.feedback is None
+    assert calls == 1
+    stored = get_context().tracing_database.get(trace_id)
+    if do_not_persist:
+        assert stored is None
+    else:
+        assert stored is not None
+        assert stored.output == [[1]]
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+def test_cached_predictions_follow_http_feedback_and_deletion(
+    isolated_context: None, is_async: bool
+) -> None:
+    calls = 0
+
+    def predict(value: int) -> int:
+        nonlocal calls
+        calls += 1
+        return value * 2
+
+    async def predict_async(value: int) -> int:
+        return predict(value)
+
+    service: Any = GreatAI.create(predict_async if is_async else predict)
+    with TestClient(service.app) as client:
+        first = client.post("/predict", json={"value": 2}).json()
+        trace_url = f"/traces/{first['trace_id']}"
+        feedback_url = f"{trace_url}/feedback/"
+        assert (
+            client.put(
+                feedback_url, json={"feedback": {"label": "correct"}}
+            ).status_code
+            == 202
+        )
+        second = client.post("/predict", json={"value": 2}).json()
+        assert second["trace_id"] == first["trace_id"]
+        assert second["feedback"] == {"label": "correct"}
+        assert second == client.get(trace_url).json()
+        assert calls == 1
+        assert client.delete(feedback_url).status_code == 204
+        assert client.post("/predict", json={"value": 2}).json()["feedback"] is None
+        assert client.delete(trace_url).status_code == 204
+        third = client.post("/predict", json={"value": 2}).json()
+        assert third["trace_id"] != first["trace_id"]
+        assert third["output"] == 4
+        assert calls == 2
+        assert client.get(f"/traces/{third['trace_id']}").json() == third
+        assert (
+            client.put(
+                f"/traces/{third['trace_id']}/feedback/", json={"feedback": "correct"}
+            ).status_code
+            == 202
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", [False, True])
+async def test_cached_predictions_observe_direct_database_changes(
+    isolated_context: None, is_async: bool
+) -> None:
+    def predict(value: int) -> int:
+        return value
+
+    async def predict_async(value: int) -> int:
+        return value
+
+    service: Any = GreatAI.create(predict_async if is_async else predict)
+    first = service(1)
+    if is_async:
+        first = await first
+    database = get_context().tracing_database
+    stored = database.get(first.trace_id)
+    assert stored is not None
+    stored.feedback = ["updated outside the service"]
+    database.update(first.trace_id, stored)
+    second = service(1)
+    if is_async:
+        second = await second
+    assert second.feedback == stored.feedback
+    second.feedback.append("caller change")
+    database.delete(first.trace_id)
+    third = service(1)
+    if is_async:
+        third = await third
+    assert third.trace_id != first.trace_id
+    assert database.get(third.trace_id) is not None