mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-16 00:14:31 +00:00
feat: durability and idempotency keys (#2450)
* feat: durability and idempotency keys * feat: more durable frontend * fix: tests * fix: mini issues * fix: better json validation * fix: tests
This commit is contained in:
112 files changed
+13520
-614
No files matched your search
@@ -710,6 +710,91 @@ class TestUpdateAgent:
|
||||
response = UpdateAgent().put(str(agent["id"]))
|
||||
assert response.status_code == 400
|
||||
|
||||
def test_publish_with_default_source_succeeds(self, app, pg_conn):
|
||||
"""The frontend's auto-selected "Default" source has no UUID — it
|
||||
sends ``source: ''`` and ``sources: []`` with ``retriever: 'classic'``.
|
||||
Pre-fix, the publish gate rejected this with
|
||||
``Missing or invalid required fields: Source``. The retriever
|
||||
carries the runtime identity, so the gate now accepts it.
|
||||
"""
|
||||
from application.api.user.agents.routes import UpdateAgent
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.prompts import PromptsRepository
|
||||
|
||||
user = "u-publish-default"
|
||||
# Draft agent with retriever='classic', no source_id, prompt + chunks set.
|
||||
agent = _seed_agent(
|
||||
pg_conn, user=user, status="draft",
|
||||
with_source=False, retriever="classic",
|
||||
)
|
||||
prompt = PromptsRepository(pg_conn).create(user, "p", "c")
|
||||
AgentsRepository(pg_conn).update(
|
||||
str(agent["id"]), user,
|
||||
{
|
||||
"prompt_id": str(prompt["id"]),
|
||||
"chunks": 2,
|
||||
"agent_type": "classic",
|
||||
},
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/update_agent/{agent['id']}", method="PUT",
|
||||
json={
|
||||
"name": "n",
|
||||
"description": "d",
|
||||
"status": "published",
|
||||
"source": "",
|
||||
"sources": [],
|
||||
"retriever": "classic",
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateAgent().put(str(agent["id"]))
|
||||
assert response.status_code == 200, (
|
||||
f"unexpected error: {response.json}"
|
||||
)
|
||||
|
||||
def test_publish_without_source_or_retriever_returns_400(
|
||||
self, app, pg_conn,
|
||||
):
|
||||
"""If neither a source nor a retriever is configured, the gate
|
||||
still trips — the agent has no way to retrieve anything."""
|
||||
from application.api.user.agents.routes import UpdateAgent
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
from application.storage.db.repositories.prompts import PromptsRepository
|
||||
|
||||
user = "u-publish-no-retriever"
|
||||
agent = _seed_agent(
|
||||
pg_conn, user=user, status="draft",
|
||||
with_source=False, retriever="",
|
||||
)
|
||||
prompt = PromptsRepository(pg_conn).create(user, "p", "c")
|
||||
AgentsRepository(pg_conn).update(
|
||||
str(agent["id"]), user,
|
||||
{
|
||||
"prompt_id": str(prompt["id"]),
|
||||
"chunks": 2,
|
||||
"agent_type": "classic",
|
||||
},
|
||||
)
|
||||
|
||||
with _patch_db(pg_conn), app.test_request_context(
|
||||
f"/api/update_agent/{agent['id']}", method="PUT",
|
||||
json={
|
||||
"name": "n",
|
||||
"description": "d",
|
||||
"status": "published",
|
||||
"source": "",
|
||||
"sources": [],
|
||||
},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = UpdateAgent().put(str(agent["id"]))
|
||||
assert response.status_code == 400
|
||||
assert "Source or retriever" in response.json.get("message", "")
|
||||
|
||||
def test_publishing_generates_api_key(self, app, pg_conn):
|
||||
from application.api.user.agents.routes import UpdateAgent
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
@@ -55,7 +55,7 @@ class TestAgentWebhookListenerGaps:
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook"
|
||||
) as mock_process:
|
||||
mock_process.delay.return_value = mock_task
|
||||
mock_process.apply_async.return_value = mock_task
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tok",
|
||||
method="POST",
|
||||
@@ -77,7 +77,7 @@ class TestAgentWebhookListenerGaps:
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook"
|
||||
) as mock_process:
|
||||
mock_process.delay.return_value = mock_task
|
||||
mock_process.apply_async.return_value = mock_task
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tok",
|
||||
method="GET",
|
||||
@@ -90,8 +90,9 @@ class TestAgentWebhookListenerGaps:
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
call_kwargs = mock_process.delay.call_args[1]
|
||||
assert call_kwargs["payload"] == {}
|
||||
call_kwargs = mock_process.apply_async.call_args[1]
|
||||
# apply_async wraps task args under ``kwargs=``.
|
||||
assert call_kwargs["kwargs"]["payload"] == {}
|
||||
|
||||
def test_enqueue_returns_task_id_in_response(self, app):
|
||||
"""Success response body includes task_id from the Celery task."""
|
||||
@@ -103,7 +104,7 @@ class TestAgentWebhookListenerGaps:
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook"
|
||||
) as mock_process:
|
||||
mock_process.delay.return_value = mock_task
|
||||
mock_process.apply_async.return_value = mock_task
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tok",
|
||||
method="POST",
|
||||
@@ -122,7 +123,7 @@ class TestAgentWebhookListenerGaps:
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook"
|
||||
) as mock_process:
|
||||
mock_process.delay.side_effect = RuntimeError("celery is down")
|
||||
mock_process.apply_async.side_effect = RuntimeError("celery is down")
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tok",
|
||||
method="POST",
|
||||
@@ -254,7 +255,7 @@ class TestAgentWebhookListener:
|
||||
fake_task = MagicMock(id="task-post")
|
||||
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.delay",
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-enq", method="POST",
|
||||
@@ -282,7 +283,7 @@ class TestAgentWebhookListener:
|
||||
fake_task = MagicMock(id="task-get")
|
||||
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.delay",
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-get?foo=bar&baz=42"
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
"""Idempotency-Key behavior on the agent webhook listener route."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.db_session", _yield
|
||||
), patch(
|
||||
"application.api.user.agents.webhooks.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _seed_agent(pg_conn, user="u", token="tk", **kw):
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
return AgentsRepository(pg_conn).create(
|
||||
user, "a", "published", incoming_webhook_token=token, **kw,
|
||||
)
|
||||
|
||||
|
||||
def _apply_async_mock():
|
||||
"""Mock for ``process_agent_webhook.apply_async``; ``task.id`` mirrors the predetermined id."""
|
||||
def _side_effect(*args, **kwargs):
|
||||
return MagicMock(id=kwargs.get("task_id") or "auto-task-id")
|
||||
return MagicMock(side_effect=_side_effect)
|
||||
|
||||
|
||||
class TestWebhookIdempotency:
|
||||
def test_no_header_enqueues_normally(self, app, pg_conn):
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-noh", token="tk-noh")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-noh", method="POST",
|
||||
json={"event": "x"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
response = listener.post(
|
||||
webhook_token="tk-noh",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_header_first_post_records_row(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-first", token="tk-first")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-first", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "key-abc"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
response = listener.post(
|
||||
webhook_token="tk-first",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
predetermined_id = apply_mock.call_args.kwargs["task_id"]
|
||||
assert response.json["task_id"] == predetermined_id
|
||||
|
||||
# Stored under the *scoped* form ``"{agent_id}:{key}"`` so two
|
||||
# agents sharing the same raw header don't collapse on PK.
|
||||
scoped_key = f"{agent['id']}:key-abc"
|
||||
row = pg_conn.execute(
|
||||
text("SELECT task_id, agent_id FROM webhook_dedup WHERE idempotency_key = :k"),
|
||||
{"k": scoped_key},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == predetermined_id
|
||||
assert str(row[1]) == str(agent["id"])
|
||||
|
||||
def test_header_forwards_idempotency_key_to_delay(self, app, pg_conn):
|
||||
"""The Celery task body needs the key so ``with_idempotency`` can
|
||||
record terminal status and ``_derive_source_id`` can pick it up.
|
||||
"""
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-fwd", token="tk-fwd")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-fwd", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "key-fwd"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
listener.post(
|
||||
webhook_token="tk-fwd",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
# Worker sees the agent-scoped form so its dedup row is also
|
||||
# agent-distinct.
|
||||
scoped_key = f"{agent['id']}:key-fwd"
|
||||
assert (
|
||||
apply_mock.call_args.kwargs["kwargs"]["idempotency_key"]
|
||||
== scoped_key
|
||||
)
|
||||
|
||||
def test_same_header_second_post_returns_cached(self, app, pg_conn):
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-rep", token="tk-rep")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tk-rep", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "key-rep"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
first = listener.post(
|
||||
webhook_token="tk-rep",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tk-rep", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "key-rep"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
second = listener.post(
|
||||
webhook_token="tk-rep",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
assert first.json == second.json
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_concurrent_same_key_only_one_apply_async(self, app, pg_engine):
|
||||
"""Race test (M3): N parallel webhook POSTs with same key → only ONE apply_async.
|
||||
|
||||
Uses ``pg_engine`` so each thread checks out its own DB connection
|
||||
(sharing a single Connection serializes at the driver level).
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
from application.storage.db.repositories.agents import AgentsRepository
|
||||
|
||||
with pg_engine.begin() as conn:
|
||||
agent = AgentsRepository(conn).create(
|
||||
"u-race", "a", "published", incoming_webhook_token="tk-race",
|
||||
)
|
||||
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
@contextmanager
|
||||
def _engine_session():
|
||||
with pg_engine.begin() as conn:
|
||||
yield conn
|
||||
|
||||
@contextmanager
|
||||
def _engine_readonly():
|
||||
with pg_engine.connect() as conn:
|
||||
yield conn
|
||||
|
||||
def fire(idx):
|
||||
# Patches sit outside the thread pool (see below); only the
|
||||
# per-thread Flask request context is set up inside.
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tk-race", method="POST",
|
||||
json={"event": idx},
|
||||
headers={"Idempotency-Key": "wh-race"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
return listener.post(
|
||||
webhook_token="tk-race",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
|
||||
# ``unittest.mock.patch`` is not thread-safe; set up
|
||||
# module-attribute patches once before fanning out so every
|
||||
# thread sees the mock instead of racing on save/restore.
|
||||
with patch(
|
||||
"application.api.user.agents.webhooks.db_session",
|
||||
_engine_session,
|
||||
), patch(
|
||||
"application.api.user.agents.webhooks.db_readonly",
|
||||
_engine_readonly,
|
||||
), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), ThreadPoolExecutor(max_workers=8) as ex:
|
||||
responses = list(ex.map(fire, range(8)))
|
||||
assert all(r.status_code == 200 for r in responses)
|
||||
assert apply_mock.call_count == 1
|
||||
ids = {r.json["task_id"] for r in responses}
|
||||
assert len(ids) == 1
|
||||
assert "deduplicated" not in ids
|
||||
|
||||
def test_same_key_different_agent_does_not_collide(self, app, pg_conn):
|
||||
"""Idempotency keys are now scoped by ``agent_id`` — two agents
|
||||
sending the same raw header each get their own dedup row, both
|
||||
requests enqueue work, and the responses carry distinct
|
||||
``task_id``s. (Pre-fix, the second agent's request was silently
|
||||
deduplicated against the first agent's row.)
|
||||
"""
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent_a = _seed_agent(pg_conn, user="u-a", token="tk-a")
|
||||
agent_b = _seed_agent(pg_conn, user="u-b", token="tk-b")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tk-a", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "global-key"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
first = listener.post(
|
||||
webhook_token="tk-a",
|
||||
agent=agent_a,
|
||||
agent_id_str=str(agent_a["id"]),
|
||||
)
|
||||
with app.test_request_context(
|
||||
"/api/webhooks/agents/tk-b", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "global-key"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
second = listener.post(
|
||||
webhook_token="tk-b",
|
||||
agent=agent_b,
|
||||
agent_id_str=str(agent_b["id"]),
|
||||
)
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
assert first.json["task_id"] != second.json["task_id"]
|
||||
assert apply_mock.call_count == 2
|
||||
|
||||
# And there are two ``webhook_dedup`` rows: one per agent scope.
|
||||
rows = pg_conn.execute(
|
||||
sql_text(
|
||||
"SELECT idempotency_key, agent_id FROM webhook_dedup "
|
||||
"WHERE idempotency_key LIKE :pat ORDER BY idempotency_key"
|
||||
),
|
||||
{"pat": "%:global-key"},
|
||||
).fetchall()
|
||||
assert len(rows) == 2
|
||||
scopes = {str(r[1]) for r in rows}
|
||||
assert scopes == {str(agent_a["id"]), str(agent_b["id"])}
|
||||
|
||||
def test_empty_header_treated_as_absent(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-empty", token="tk-empty")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-empty", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": ""},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
response = listener.post(
|
||||
webhook_token="tk-empty",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
count = pg_conn.execute(
|
||||
text("SELECT count(*) FROM webhook_dedup")
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
def test_oversized_header_rejected_with_400(self, app, pg_conn):
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-big", token="tk-big")
|
||||
oversized = "x" * 257
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
) as mock_apply, app.test_request_context(
|
||||
"/api/webhooks/agents/tk-big", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": oversized},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
response = listener.post(
|
||||
webhook_token="tk-big",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert mock_apply.call_count == 0
|
||||
|
||||
def test_stale_dedup_row_does_not_block_new_work(self, app, pg_conn):
|
||||
"""Regression for the TTL fail-shut bug: a >24h-old dedup row
|
||||
must not silently drop a new request. Pre-fix, the second POST
|
||||
returned ``task_id="deduplicated"`` and never enqueued.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.agents.webhooks import AgentWebhookListener
|
||||
|
||||
agent = _seed_agent(pg_conn, user="u-stale", token="tk-stale")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
# First POST creates a dedup row.
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-stale", method="POST",
|
||||
json={"event": "x"},
|
||||
headers={"Idempotency-Key": "stale-key"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
first = listener.post(
|
||||
webhook_token="tk-stale",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert first.status_code == 200
|
||||
first_task_id = first.json["task_id"]
|
||||
assert first_task_id != "deduplicated"
|
||||
|
||||
# Backdate the row so it looks 25h old.
|
||||
scoped_key = f"{agent['id']}:stale-key"
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE webhook_dedup SET created_at = "
|
||||
"clock_timestamp() - make_interval(hours => 25) "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": scoped_key},
|
||||
)
|
||||
|
||||
# Second POST with the same key must enqueue again, not silently dedup.
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.agents.webhooks.process_agent_webhook.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/webhooks/agents/tk-stale", method="POST",
|
||||
json={"event": "x2"},
|
||||
headers={"Idempotency-Key": "stale-key"},
|
||||
):
|
||||
listener = AgentWebhookListener()
|
||||
second = listener.post(
|
||||
webhook_token="tk-stale",
|
||||
agent=agent,
|
||||
agent_id_str=str(agent["id"]),
|
||||
)
|
||||
assert second.status_code == 200
|
||||
assert second.json["task_id"] != "deduplicated"
|
||||
assert second.json["task_id"] != first_task_id
|
||||
assert apply_mock.call_count == 2
|
||||
@@ -107,7 +107,7 @@ class TestUploadFile:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.delay",
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
@@ -164,7 +164,7 @@ class TestUploadFile:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.delay",
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
@@ -197,7 +197,7 @@ class TestUploadFile:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.delay",
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
return_value=MagicMock(id="t"),
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
@@ -274,7 +274,7 @@ class TestUploadRemote:
|
||||
|
||||
fake_task = MagicMock(id="remote-task-1")
|
||||
with patch(
|
||||
"application.api.user.sources.upload.ingest_remote.delay",
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
@@ -295,7 +295,7 @@ class TestUploadRemote:
|
||||
|
||||
fake_task = MagicMock(id="url-task")
|
||||
with patch(
|
||||
"application.api.user.sources.upload.ingest_remote.delay",
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
@@ -315,7 +315,7 @@ class TestUploadRemote:
|
||||
|
||||
fake_task = MagicMock(id="reddit-task")
|
||||
with patch(
|
||||
"application.api.user.sources.upload.ingest_remote.delay",
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
@@ -334,7 +334,7 @@ class TestUploadRemote:
|
||||
from application.api.user.sources.upload import UploadRemote
|
||||
|
||||
with patch(
|
||||
"application.api.user.sources.upload.ingest_remote.delay",
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
side_effect=RuntimeError("boom"),
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
@@ -460,7 +460,7 @@ class TestManageSourceFiles:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.delay",
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/manage_source_files", method="POST",
|
||||
@@ -560,7 +560,7 @@ class TestManageSourceFiles:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.delay",
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/manage_source_files", method="POST",
|
||||
@@ -666,7 +666,7 @@ class TestManageSourceFiles:
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.delay",
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
return_value=fake_task,
|
||||
), app.test_request_context(
|
||||
"/api/manage_source_files", method="POST",
|
||||
|
||||
@@ -0,0 +1,947 @@
|
||||
"""Idempotency-Key behavior on the /api/upload and /api/remote routes."""
|
||||
|
||||
import io
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.sources.upload.db_session", _yield
|
||||
), patch(
|
||||
"application.api.user.sources.upload.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _apply_async_mock():
|
||||
"""Mock for ``ingest.apply_async``; ``task.id`` mirrors the predetermined id."""
|
||||
def _side_effect(*args, **kwargs):
|
||||
return MagicMock(id=kwargs.get("task_id") or "auto-task-id")
|
||||
m = MagicMock(side_effect=_side_effect)
|
||||
return m
|
||||
|
||||
|
||||
class TestUploadIdempotency:
|
||||
def test_no_header_enqueues_normally(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
response = UploadFile().post()
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
# No key → no predetermined id was passed.
|
||||
assert "task_id" not in apply_mock.call_args.kwargs
|
||||
|
||||
def test_header_first_post_records_row(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "up-key-1"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
response = UploadFile().post()
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
predetermined_id = apply_mock.call_args.kwargs["task_id"]
|
||||
assert response.json["task_id"] == predetermined_id
|
||||
|
||||
# The dedup row is keyed on the *scoped* form ``"{user}:{key}"``
|
||||
# so two users sending the same raw header don't collapse.
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT task_id, task_name, status FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "alice:up-key-1"},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == predetermined_id
|
||||
assert row[1] == "ingest"
|
||||
assert row[2] == "pending"
|
||||
|
||||
def test_header_forwards_idempotency_key_to_delay(self, app, pg_conn):
|
||||
"""The Celery task body needs the key so ``with_idempotency`` can
|
||||
record terminal status and ``_derive_source_id`` can pick it up.
|
||||
"""
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"c"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "up-fwd"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
UploadFile().post()
|
||||
# Worker sees the scoped form so its ``with_idempotency`` row
|
||||
# and ``_derive_source_id`` are also user-distinct.
|
||||
assert (
|
||||
apply_mock.call_args.kwargs["kwargs"]["idempotency_key"]
|
||||
== "alice:up-fwd"
|
||||
)
|
||||
|
||||
def test_same_header_second_post_returns_cached(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "up-rep"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
first = UploadFile().post()
|
||||
with app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "up-rep"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
second = UploadFile().post()
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
assert first.json == second.json
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_same_key_different_users_does_not_collide(self, app, pg_conn):
|
||||
"""Cross-user collision regression: two users sending the same
|
||||
raw ``Idempotency-Key`` must each get their own dedup row, both
|
||||
requests enqueue, and the responses carry distinct ``task_id``s.
|
||||
(Pre-fix, the second user's request was silently deduplicated
|
||||
against the first user's row.)
|
||||
"""
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
def _fire(user):
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": user, "name": "j",
|
||||
"file": (io.BytesIO(b"c"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "shared-key"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
return UploadFile().post()
|
||||
|
||||
first = _fire("alice")
|
||||
second = _fire("bob")
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
assert first.json["task_id"] != second.json["task_id"]
|
||||
assert apply_mock.call_count == 2
|
||||
|
||||
rows = pg_conn.execute(
|
||||
sql_text(
|
||||
"SELECT idempotency_key FROM task_dedup "
|
||||
"WHERE idempotency_key LIKE :pat ORDER BY idempotency_key"
|
||||
),
|
||||
{"pat": "%:shared-key"},
|
||||
).fetchall()
|
||||
assert {r[0] for r in rows} == {"alice:shared-key", "bob:shared-key"}
|
||||
|
||||
def test_concurrent_same_key_only_one_apply_async(self, app, pg_engine):
|
||||
"""Race test (M3): N parallel POSTs with same key → only ONE apply_async.
|
||||
|
||||
Uses ``pg_engine`` (not ``pg_conn``) so each thread can check out
|
||||
its own DB connection — sharing a single Connection across
|
||||
threads serializes at the driver level and defeats the race.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
@contextmanager
|
||||
def _engine_session():
|
||||
with pg_engine.begin() as conn:
|
||||
yield conn
|
||||
|
||||
@contextmanager
|
||||
def _engine_readonly():
|
||||
with pg_engine.connect() as conn:
|
||||
yield conn
|
||||
|
||||
def fire(idx):
|
||||
# Patches sit OUTSIDE the threads (see below); only the
|
||||
# per-thread Flask request context is set up inside.
|
||||
with app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"x"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "up-race"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
return UploadFile().post()
|
||||
|
||||
# ``unittest.mock.patch`` is not thread-safe — concurrent
|
||||
# ``__enter__`` calls race on saving/restoring the module
|
||||
# attribute and can leave threads pointing at the real
|
||||
# function instead of the mock. Set up patches once, share
|
||||
# across threads.
|
||||
with patch(
|
||||
"application.api.user.sources.upload.db_session",
|
||||
_engine_session,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.db_readonly",
|
||||
_engine_readonly,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), ThreadPoolExecutor(max_workers=8) as ex:
|
||||
responses = list(ex.map(fire, range(8)))
|
||||
assert all(r.status_code == 200 for r in responses)
|
||||
# Only one writer wins the claim, so only one apply_async is fired.
|
||||
assert apply_mock.call_count == 1
|
||||
# All 8 responses share the same task_id (winner's predetermined id).
|
||||
ids = {r.json["task_id"] for r in responses}
|
||||
assert len(ids) == 1
|
||||
assert "deduplicated" not in ids
|
||||
|
||||
def test_empty_header_treated_as_absent(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": ""},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
response = UploadFile().post()
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
count = pg_conn.execute(
|
||||
text("SELECT count(*) FROM task_dedup")
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
def test_oversized_header_rejected_with_400(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
oversized = "x" * 257
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
) as mock_apply, app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": oversized},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
response = UploadFile().post()
|
||||
assert response.status_code == 400
|
||||
assert mock_apply.call_count == 0
|
||||
|
||||
def test_stale_dedup_row_does_not_block_new_work(self, app, pg_conn):
|
||||
"""Regression for the TTL fail-shut bug: a >24h-old dedup row
|
||||
must not silently drop a new upload. Pre-fix, the second POST
|
||||
returned ``task_id="deduplicated"`` and never enqueued.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import UploadFile
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.ingest.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "stale-up"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
first = UploadFile().post()
|
||||
assert first.status_code == 200
|
||||
first_task_id = first.json["task_id"]
|
||||
assert first_task_id != "deduplicated"
|
||||
|
||||
# Backdate the row past TTL.
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE task_dedup SET created_at = "
|
||||
"clock_timestamp() - make_interval(hours => 25) "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "alice:stale-up"},
|
||||
)
|
||||
|
||||
with app.test_request_context(
|
||||
"/api/upload", method="POST",
|
||||
data={
|
||||
"user": "alice", "name": "j",
|
||||
"file": (io.BytesIO(b"content"), "doc.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "stale-up"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
second = UploadFile().post()
|
||||
|
||||
assert second.status_code == 200
|
||||
assert second.json["task_id"] != "deduplicated"
|
||||
assert second.json["task_id"] != first_task_id
|
||||
assert apply_mock.call_count == 2
|
||||
|
||||
|
||||
class TestRemoteIdempotency:
|
||||
def test_no_header_enqueues_normally(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadRemote
|
||||
|
||||
apply_mock = _apply_async_mock()
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
data={
|
||||
"user": "u", "source": "github", "name": "g",
|
||||
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u"}
|
||||
response = UploadRemote().post()
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_header_first_post_records_row(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import UploadRemote
|
||||
|
||||
apply_mock = _apply_async_mock()
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
data={
|
||||
"user": "u", "source": "github", "name": "g",
|
||||
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "rem-key-1"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u"}
|
||||
response = UploadRemote().post()
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
predetermined_id = apply_mock.call_args.kwargs["task_id"]
|
||||
# Scoped key: ``"{user}:{key}"``.
|
||||
row = pg_conn.execute(
|
||||
text("SELECT task_id, task_name FROM task_dedup WHERE idempotency_key = :k"),
|
||||
{"k": "u:rem-key-1"},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == predetermined_id
|
||||
assert row[1] == "ingest_remote"
|
||||
|
||||
def test_same_header_second_post_returns_cached(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadRemote
|
||||
|
||||
apply_mock = _apply_async_mock()
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
data={
|
||||
"user": "u", "source": "github", "name": "g",
|
||||
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "rem-rep"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u"}
|
||||
first = UploadRemote().post()
|
||||
with app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
data={
|
||||
"user": "u", "source": "github", "name": "g",
|
||||
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "rem-rep"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u"}
|
||||
second = UploadRemote().post()
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
assert first.json == second.json
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_oversized_header_rejected_with_400(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import UploadRemote
|
||||
|
||||
oversized = "x" * 257
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.ingest_remote.apply_async",
|
||||
) as mock_apply, app.test_request_context(
|
||||
"/api/remote", method="POST",
|
||||
data={
|
||||
"user": "u", "source": "github", "name": "g",
|
||||
"data": json.dumps({"repo_url": "https://github.com/x/y"}),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": oversized},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": "u"}
|
||||
response = UploadRemote().post()
|
||||
assert response.status_code == 400
|
||||
assert mock_apply.call_count == 0
|
||||
|
||||
|
||||
def _seed_source(pg_conn, user="u", **kw):
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
return SourcesRepository(pg_conn).create("manage-src", user_id=user, **kw)
|
||||
|
||||
|
||||
class TestManageSourceFilesIdempotency:
|
||||
"""Same-key dedup contract for ``ManageSourceFiles.post``: a duplicate
|
||||
POST must not enqueue a second ``reingest_source_task``. The worker
|
||||
decorator only deduplicates *post-completion*, so the HTTP handler is
|
||||
the only place that can serialize concurrent in-flight requests.
|
||||
"""
|
||||
|
||||
def _add_request(self, app, src_id, user, key=None):
|
||||
kwargs = dict(
|
||||
data={
|
||||
"source_id": str(src_id),
|
||||
"operation": "add",
|
||||
"file": (io.BytesIO(b"content"), "new.txt"),
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
if key is not None:
|
||||
kwargs["headers"] = {"Idempotency-Key": key}
|
||||
return app.test_request_context(
|
||||
"/api/manage_source_files", method="POST", **kwargs,
|
||||
)
|
||||
|
||||
def test_no_header_enqueues_normally_no_claim_row(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-noh"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), self._add_request(app, src["id"], user):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
# No key → predetermined task_id is None, Celery generates one.
|
||||
assert apply_mock.call_args.kwargs["task_id"] is None
|
||||
n = pg_conn.execute(
|
||||
text("SELECT count(*) FROM task_dedup")
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_header_records_dedup_row_with_predetermined_id(self, app, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-rec"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), self._add_request(app, src["id"], user, key="mgr-key-1"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert apply_mock.call_count == 1
|
||||
predetermined_id = apply_mock.call_args.kwargs["task_id"]
|
||||
assert predetermined_id is not None
|
||||
assert response.json["reingest_task_id"] == predetermined_id
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT task_id, task_name, status FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": f"{user}:mgr-key-1"},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == predetermined_id
|
||||
assert row[1] == "reingest_source_task"
|
||||
assert row[2] == "pending"
|
||||
|
||||
def test_same_key_second_post_returns_cached(self, app, pg_conn):
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-rep"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
):
|
||||
with self._add_request(app, src["id"], user, key="mgr-rep"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
first = ManageSourceFiles().post()
|
||||
with self._add_request(app, src["id"], user, key="mgr-rep"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
second = ManageSourceFiles().post()
|
||||
|
||||
assert first.status_code == 200
|
||||
assert second.status_code == 200
|
||||
# Loser short-circuits before any storage mutation.
|
||||
assert apply_mock.call_count == 1
|
||||
# Loser's response carries the winner's task_id, not the
|
||||
# original 200-with-added_files payload.
|
||||
assert second.json["task_id"] == first.json["reingest_task_id"]
|
||||
# Confirm the loser never invoked the file-save path.
|
||||
assert fake_storage.save_file.call_count == 1
|
||||
|
||||
def test_concurrent_same_key_only_one_apply_async(self, app, pg_engine):
|
||||
"""N parallel same-key POSTs → exactly one apply_async."""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
from application.storage.db.repositories.sources import (
|
||||
SourcesRepository,
|
||||
)
|
||||
|
||||
user = "alice-mgr-race"
|
||||
with pg_engine.begin() as conn:
|
||||
src = SourcesRepository(conn).create(
|
||||
"race-src", user_id=user, file_path="/data",
|
||||
)
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
@contextmanager
|
||||
def _engine_session():
|
||||
with pg_engine.begin() as conn:
|
||||
yield conn
|
||||
|
||||
@contextmanager
|
||||
def _engine_readonly():
|
||||
with pg_engine.connect() as conn:
|
||||
yield conn
|
||||
|
||||
def fire(_idx):
|
||||
# Patches sit outside the thread pool (see below); only the
|
||||
# per-thread Flask request context is set up inside.
|
||||
with self._add_request(app, src["id"], user, key="mgr-race"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
return ManageSourceFiles().post()
|
||||
|
||||
# ``unittest.mock.patch`` is not thread-safe; set up the
|
||||
# module-attribute patches once before fanning out so every
|
||||
# thread sees the mock instead of racing on save/restore.
|
||||
with patch(
|
||||
"application.api.user.sources.upload.db_session",
|
||||
_engine_session,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.db_readonly",
|
||||
_engine_readonly,
|
||||
), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), ThreadPoolExecutor(max_workers=8) as ex:
|
||||
responses = list(ex.map(fire, range(8)))
|
||||
assert all(r.status_code == 200 for r in responses)
|
||||
assert apply_mock.call_count == 1
|
||||
|
||||
def test_remove_directory_failure_releases_claim(self, app, pg_conn):
|
||||
"""When storage.remove_directory returns False the handler must
|
||||
release the dedup row so a client retry can win the claim. Without
|
||||
the release, the next retry would silently 200-cache to a task_id
|
||||
that was never enqueued.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-rmfail"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
fake_storage.is_directory.return_value = True
|
||||
fake_storage.remove_directory.return_value = False
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), app.test_request_context(
|
||||
"/api/manage_source_files", method="POST",
|
||||
data={
|
||||
"source_id": str(src["id"]),
|
||||
"operation": "remove_directory",
|
||||
"directory_path": "subdir",
|
||||
},
|
||||
content_type="multipart/form-data",
|
||||
headers={"Idempotency-Key": "mgr-rmfail"},
|
||||
):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 500
|
||||
assert apply_mock.call_count == 0
|
||||
# Claim row was released so a retry can re-claim.
|
||||
n = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE idempotency_key = :k AND status = 'pending'"
|
||||
),
|
||||
{"k": f"{user}:mgr-rmfail"},
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_storage_save_failure_releases_claim(self, app, pg_conn):
|
||||
"""Regression: an exception from ``storage.save_file`` after the
|
||||
claim must release the dedup row. Pre-fix the outer ``except``
|
||||
logged + 500'd without releasing, so a retry within 24h returned
|
||||
a cached predetermined ``task_id`` for a task that never enqueued.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-storefail"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
fake_storage.save_file.side_effect = RuntimeError("disk full")
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), self._add_request(app, src["id"], user, key="mgr-storefail"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 500
|
||||
assert apply_mock.call_count == 0
|
||||
n = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": f"{user}:mgr-storefail"},
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_apply_async_failure_releases_claim(self, app, pg_conn):
|
||||
"""If the broker is unreachable, ``apply_async`` raises *after*
|
||||
the claim was made. The claim row must be released so a retry
|
||||
can re-claim with a fresh task_id.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-brokerdown"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
|
||||
def _broker_down(*args, **kwargs):
|
||||
raise ConnectionError("broker unreachable")
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
side_effect=_broker_down,
|
||||
), self._add_request(app, src["id"], user, key="mgr-brokerdown"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 500
|
||||
n = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": f"{user}:mgr-brokerdown"},
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_db_update_failure_releases_claim(self, app, pg_conn):
|
||||
"""A DB blip during the ``file_name_map`` update happens after
|
||||
storage mutated and the claim was made. Must still release.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
from application.storage.db.repositories import sources as src_module
|
||||
|
||||
user = "alice-mgr-dbfail"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
original_update = src_module.SourcesRepository.update
|
||||
|
||||
def _explode(self, *args, **kwargs):
|
||||
raise RuntimeError("db transient")
|
||||
|
||||
src_module.SourcesRepository.update = _explode
|
||||
try:
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), self._add_request(app, src["id"], user, key="mgr-dbfail"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
finally:
|
||||
src_module.SourcesRepository.update = original_update
|
||||
|
||||
assert response.status_code == 500
|
||||
# apply_async never ran because the DB raise pre-empted it.
|
||||
assert apply_mock.call_count == 0
|
||||
n = pg_conn.execute(
|
||||
text(
|
||||
"SELECT count(*) FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": f"{user}:mgr-dbfail"},
|
||||
).scalar()
|
||||
assert n == 0
|
||||
|
||||
def test_successful_path_keeps_claim_for_worker(self, app, pg_conn):
|
||||
"""The claim row must persist after a successful ``apply_async``
|
||||
— the worker owns the predetermined task_id and same-key retries
|
||||
should resolve to the in-flight task, not re-enqueue.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.sources.upload import ManageSourceFiles
|
||||
|
||||
user = "alice-mgr-keep"
|
||||
src = _seed_source(pg_conn, user=user, file_path="/data")
|
||||
|
||||
fake_storage = MagicMock()
|
||||
apply_mock = _apply_async_mock()
|
||||
|
||||
with _patch_db(pg_conn), patch(
|
||||
"application.api.user.sources.upload.StorageCreator.get_storage",
|
||||
return_value=fake_storage,
|
||||
), patch(
|
||||
"application.api.user.tasks.reingest_source_task.apply_async",
|
||||
apply_mock,
|
||||
), self._add_request(app, src["id"], user, key="mgr-keep"):
|
||||
from flask import request
|
||||
request.decoded_token = {"sub": user}
|
||||
response = ManageSourceFiles().post()
|
||||
|
||||
assert response.status_code == 200
|
||||
# Row is still there in pending status — worker has not finalised yet.
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": f"{user}:mgr-keep"},
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
assert row[0] == "pending"
|
||||
@@ -419,6 +419,94 @@ class TestGetSingleConversationHappy:
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGetMessageTail:
|
||||
"""Tail-poll endpoint (``GET /api/messages/<id>/tail``) used by the
|
||||
frontend to recover a placeholder/streaming row after a refresh.
|
||||
"""
|
||||
|
||||
def _seed_in_flight_message(self, pg_conn, owner_user_id):
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
conv_id = _seed_conversation(pg_conn, owner_user_id, name="streaming chat")
|
||||
repo = ConversationsRepository(pg_conn)
|
||||
msg = repo.reserve_message(
|
||||
conv_id,
|
||||
prompt="what's happening?",
|
||||
placeholder_response=(
|
||||
"Response was terminated prior to completion, try regenerating."
|
||||
),
|
||||
request_id=str(uuid.uuid4()),
|
||||
status="streaming",
|
||||
)
|
||||
return conv_id, str(msg["id"])
|
||||
|
||||
def test_owner_can_tail(self, app, pg_conn):
|
||||
from application.api.user.conversations.routes import GetMessageTail
|
||||
|
||||
owner = "user-owner"
|
||||
_, msg_id = self._seed_in_flight_message(pg_conn, owner)
|
||||
|
||||
with _patch_conversations_db(pg_conn), app.test_request_context(
|
||||
f"/api/messages/{msg_id}/tail"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": owner}
|
||||
response = GetMessageTail().get(msg_id)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json["status"] == "streaming"
|
||||
assert response.json["message_id"] == msg_id
|
||||
|
||||
def test_shared_user_can_tail(self, app, pg_conn):
|
||||
"""A user in ``conversations.shared_with`` must be able to tail
|
||||
an in-flight placeholder. Without the shared-with predicate
|
||||
here, ``get_single_conversation`` lets them load the row but
|
||||
the tail-poll silently 404s and the in-flight bubble never
|
||||
resolves on the shared user's side.
|
||||
"""
|
||||
from application.api.user.conversations.routes import GetMessageTail
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
owner = "user-owner-shared"
|
||||
shared_user = "user-shared"
|
||||
conv_id, msg_id = self._seed_in_flight_message(pg_conn, owner)
|
||||
ConversationsRepository(pg_conn).add_shared_user(conv_id, shared_user)
|
||||
|
||||
with _patch_conversations_db(pg_conn), app.test_request_context(
|
||||
f"/api/messages/{msg_id}/tail"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": shared_user}
|
||||
response = GetMessageTail().get(msg_id)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json["message_id"] == msg_id
|
||||
|
||||
def test_non_member_gets_404(self, app, pg_conn):
|
||||
from application.api.user.conversations.routes import GetMessageTail
|
||||
|
||||
owner = "user-owner-private"
|
||||
intruder = "user-intruder"
|
||||
_, msg_id = self._seed_in_flight_message(pg_conn, owner)
|
||||
|
||||
with _patch_conversations_db(pg_conn), app.test_request_context(
|
||||
f"/api/messages/{msg_id}/tail"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": intruder}
|
||||
response = GetMessageTail().get(msg_id)
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
class TestUpdateConversationNameHappy:
|
||||
def test_returns_401_unauthenticated(self, app):
|
||||
from application.api.user.conversations.routes import (
|
||||
|
||||
@@ -0,0 +1,419 @@
|
||||
"""Unit-level behavior of the per-Celery-task idempotency wrapper."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_decorator_db(conn):
|
||||
"""Route the wrapper's ``db_session`` / ``db_readonly`` at ``conn``."""
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.idempotency.db_session", _yield
|
||||
), patch(
|
||||
"application.api.user.idempotency.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _fake_celery_self(task_id="task-123"):
|
||||
"""Minimal stand-in mirroring ``self.request.id`` on a Celery task."""
|
||||
self_ = MagicMock(name="celery_self")
|
||||
self_.request.id = task_id
|
||||
return self_
|
||||
|
||||
|
||||
def _row_for(conn, key):
|
||||
return conn.execute(
|
||||
text(
|
||||
"SELECT task_name, task_id, status, result_json "
|
||||
"FROM task_dedup WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": key},
|
||||
).fetchone()
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestNoKey:
|
||||
def test_pass_through_no_db_hit(self, pg_conn):
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
calls = []
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, x, idempotency_key=None):
|
||||
calls.append(x)
|
||||
return {"x": x}
|
||||
|
||||
with patch(
|
||||
"application.api.user.idempotency.db_session"
|
||||
) as mock_session, patch(
|
||||
"application.api.user.idempotency.db_readonly"
|
||||
) as mock_readonly:
|
||||
result = task(_fake_celery_self(), 7)
|
||||
|
||||
assert result == {"x": 7}
|
||||
assert calls == [7]
|
||||
assert mock_session.call_count == 0
|
||||
assert mock_readonly.call_count == 0
|
||||
|
||||
def test_empty_string_key_treated_as_absent(self, pg_conn):
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
return {"ran": True}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
result = task(_fake_celery_self(), idempotency_key="")
|
||||
|
||||
assert result == {"ran": True}
|
||||
count = pg_conn.execute(
|
||||
text("SELECT count(*) FROM task_dedup")
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestFirstRunWithKey:
|
||||
def test_records_completed_row(self, pg_conn):
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
return {"answer": 42}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
result = task(_fake_celery_self("tid-1"), idempotency_key="k-first")
|
||||
|
||||
assert result == {"answer": 42}
|
||||
row = _row_for(pg_conn, "k-first")
|
||||
assert row is not None
|
||||
assert row[0] == "thing"
|
||||
assert row[1] == "tid-1"
|
||||
assert row[2] == "completed"
|
||||
assert row[3] == {"answer": 42}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSecondRunCompletedShortCircuits:
|
||||
def test_returns_cached_without_invoking(self, pg_conn):
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
invocations = {"count": 0}
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, value, idempotency_key=None):
|
||||
invocations["count"] += 1
|
||||
return {"value": value, "n": invocations["count"]}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
first = task(_fake_celery_self("tid-A"), 1, idempotency_key="k-rep")
|
||||
second = task(_fake_celery_self("tid-B"), 2, idempotency_key="k-rep")
|
||||
|
||||
assert first == {"value": 1, "n": 1}
|
||||
assert second == first
|
||||
assert invocations["count"] == 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestFirstRunFails:
|
||||
def test_propagates_and_leaves_pending(self, pg_conn):
|
||||
"""An exception propagates so Celery's retry policy fires; the row
|
||||
stays in ``pending`` (with bumped attempt_count) so the next
|
||||
attempt isn't gated as already-completed.
|
||||
"""
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
raise RuntimeError("kaboom")
|
||||
|
||||
with _patch_decorator_db(pg_conn), pytest.raises(RuntimeError, match="kaboom"):
|
||||
task(_fake_celery_self("tid-X"), idempotency_key="k-fail")
|
||||
|
||||
row = _row_for(pg_conn, "k-fail")
|
||||
assert row is not None
|
||||
assert row[0] == "thing"
|
||||
assert row[2] == "pending"
|
||||
assert row[3] is None
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPoisonLoopGuard:
|
||||
def test_refuses_after_max_attempts(self, pg_conn):
|
||||
from application.api.user.idempotency import (
|
||||
MAX_TASK_ATTEMPTS, with_idempotency,
|
||||
)
|
||||
|
||||
invocations = {"count": 0}
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
invocations["count"] += 1
|
||||
raise RuntimeError("never converges")
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
for _ in range(MAX_TASK_ATTEMPTS):
|
||||
with pytest.raises(RuntimeError):
|
||||
task(_fake_celery_self(), idempotency_key="k-poison")
|
||||
# The next entry trips the guard and *does not* call fn.
|
||||
result = task(_fake_celery_self(), idempotency_key="k-poison")
|
||||
assert invocations["count"] == MAX_TASK_ATTEMPTS
|
||||
assert result["success"] is False
|
||||
assert "poison-loop" in result["error"]
|
||||
row = _row_for(pg_conn, "k-poison")
|
||||
assert row[2] == "failed"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestPreviousPendingReruns:
|
||||
def test_pending_row_does_not_short_circuit(self, pg_conn):
|
||||
"""HTTP boundary writes ``pending``; on first arrival the wrapper still runs."""
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
from application.storage.db.repositories.idempotency import (
|
||||
IdempotencyRepository,
|
||||
)
|
||||
|
||||
IdempotencyRepository(pg_conn).claim_task(
|
||||
key="k-pending-prior", task_name="thing",
|
||||
task_id="enq-task-id",
|
||||
)
|
||||
|
||||
invocations = {"count": 0}
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
invocations["count"] += 1
|
||||
return {"final": "answer"}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
result = task(
|
||||
_fake_celery_self("worker-tid"),
|
||||
idempotency_key="k-pending-prior",
|
||||
)
|
||||
|
||||
assert result == {"final": "answer"}
|
||||
assert invocations["count"] == 1
|
||||
row = _row_for(pg_conn, "k-pending-prior")
|
||||
assert row[2] == "completed"
|
||||
assert row[3] == {"final": "answer"}
|
||||
# The HTTP-claimed task_id is preserved across worker bumps so
|
||||
# losers still see the same predetermined id.
|
||||
assert row[1] == "enq-task-id"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRaceWithCompletedRow:
|
||||
"""A second worker finishing after the first should not clobber the completed row."""
|
||||
|
||||
def test_second_record_no_ops_on_completed(self, pg_conn):
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
from application.storage.db.repositories.idempotency import (
|
||||
IdempotencyRepository,
|
||||
)
|
||||
|
||||
# Seed a completed row directly via SQL (bypasses claim_task's
|
||||
# ON CONFLICT DO NOTHING guard so this row exists for the test).
|
||||
from sqlalchemy import text
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"INSERT INTO task_dedup (idempotency_key, task_name, task_id, "
|
||||
"result_json, status) VALUES (:k, :tn, :tid, "
|
||||
"CAST(:rj AS jsonb), 'completed')"
|
||||
),
|
||||
{
|
||||
"k": "k-race", "tn": "thing", "tid": "winner",
|
||||
"rj": '{"who": "winner"}',
|
||||
},
|
||||
)
|
||||
# IdempotencyRepository present but unused — keep as smoke check
|
||||
# the constructor still works.
|
||||
IdempotencyRepository(pg_conn)
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
return {"who": "loser"}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
result = task(_fake_celery_self("loser-tid"), idempotency_key="k-race")
|
||||
|
||||
assert result == {"who": "winner"}
|
||||
row = _row_for(pg_conn, "k-race")
|
||||
assert row[1] == "winner"
|
||||
assert row[3] == {"who": "winner"}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLiveLeaseDefersConcurrentRun:
|
||||
"""Two-worker concurrency: Worker 1 owns a fresh lease, Worker 2's
|
||||
redelivery must ``self.retry`` instead of running the task body.
|
||||
"""
|
||||
|
||||
def test_second_worker_reraises_retry_without_running(self, pg_conn):
|
||||
from application.api.user.idempotency import (
|
||||
LEASE_TTL_SECONDS, with_idempotency,
|
||||
)
|
||||
from application.storage.db.repositories.idempotency import (
|
||||
IdempotencyRepository,
|
||||
)
|
||||
|
||||
# Worker 1 has already claimed a fresh lease.
|
||||
IdempotencyRepository(pg_conn).try_claim_lease(
|
||||
key="k-busy", task_name="thing",
|
||||
task_id="t-worker-1", owner_id="worker-1",
|
||||
)
|
||||
|
||||
invocations = {"count": 0}
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
invocations["count"] += 1
|
||||
return {"ran": True}
|
||||
|
||||
# ``self.retry`` raises celery.exceptions.Retry; mirror that
|
||||
# contract here so the wrapper's ``raise self.retry(...)``
|
||||
# propagates a real exception.
|
||||
class _RetrySignal(Exception):
|
||||
pass
|
||||
|
||||
worker2 = _fake_celery_self("t-worker-2")
|
||||
worker2.retry.side_effect = _RetrySignal("retry scheduled")
|
||||
|
||||
with _patch_decorator_db(pg_conn), pytest.raises(_RetrySignal):
|
||||
task(worker2, idempotency_key="k-busy")
|
||||
|
||||
# Task body never executed.
|
||||
assert invocations["count"] == 0
|
||||
# ``self.retry`` was invoked with countdown == lease TTL.
|
||||
worker2.retry.assert_called_once()
|
||||
kwargs = worker2.retry.call_args.kwargs
|
||||
assert kwargs.get("countdown") == LEASE_TTL_SECONDS
|
||||
|
||||
def test_expired_lease_can_be_reclaimed_by_next_worker(self, pg_conn):
|
||||
"""After the lease TTL elapses, the next attempt claims and runs
|
||||
— the original "Worker 1 died" recovery path.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
from application.storage.db.repositories.idempotency import (
|
||||
IdempotencyRepository,
|
||||
)
|
||||
|
||||
IdempotencyRepository(pg_conn).try_claim_lease(
|
||||
key="k-stale", task_name="thing",
|
||||
task_id="t-1", owner_id="dead-worker",
|
||||
)
|
||||
# Force the lease into the past.
|
||||
pg_conn.execute(
|
||||
text(
|
||||
"UPDATE task_dedup "
|
||||
"SET lease_expires_at = clock_timestamp() "
|
||||
" - make_interval(secs => 5) "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "k-stale"},
|
||||
)
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
return {"ran": True}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
result = task(_fake_celery_self("t-2"), idempotency_key="k-stale")
|
||||
|
||||
assert result == {"ran": True}
|
||||
row = _row_for(pg_conn, "k-stale")
|
||||
assert row[2] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestExceptionPathReleasesLease:
|
||||
"""When ``fn`` raises, the lease is dropped so the next attempt
|
||||
doesn't have to wait the full TTL before re-claiming.
|
||||
"""
|
||||
|
||||
def test_release_clears_lease_owner(self, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with _patch_decorator_db(pg_conn), pytest.raises(RuntimeError):
|
||||
task(_fake_celery_self("tid-1"), idempotency_key="k-released")
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, lease_owner_id, lease_expires_at "
|
||||
"FROM task_dedup WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "k-released"},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
assert row[1] is None
|
||||
assert row[2] is None
|
||||
|
||||
def test_next_attempt_can_reclaim_after_release(self, pg_conn):
|
||||
"""Sequential retries don't get blocked by the lease TTL."""
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
invocations = {"count": 0}
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
invocations["count"] += 1
|
||||
if invocations["count"] < 3:
|
||||
raise RuntimeError("transient")
|
||||
return {"finally": True}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
for _ in range(2):
|
||||
with pytest.raises(RuntimeError):
|
||||
task(_fake_celery_self(), idempotency_key="k-converge")
|
||||
result = task(_fake_celery_self(), idempotency_key="k-converge")
|
||||
|
||||
assert invocations["count"] == 3
|
||||
assert result == {"finally": True}
|
||||
row = _row_for(pg_conn, "k-converge")
|
||||
assert row[2] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestSuccessfulRunClearsLease:
|
||||
"""``finalize_task`` clears the lease columns so operator dashboards
|
||||
don't show stale ``lease_expires_at`` on completed rows.
|
||||
"""
|
||||
|
||||
def test_completed_row_has_null_lease(self, pg_conn):
|
||||
from sqlalchemy import text
|
||||
|
||||
from application.api.user.idempotency import with_idempotency
|
||||
|
||||
@with_idempotency(task_name="thing")
|
||||
def task(self, idempotency_key=None):
|
||||
return {"ok": True}
|
||||
|
||||
with _patch_decorator_db(pg_conn):
|
||||
task(_fake_celery_self("tid-1"), idempotency_key="k-done")
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, lease_owner_id, lease_expires_at "
|
||||
"FROM task_dedup WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "k-done"},
|
||||
).fetchone()
|
||||
assert row[0] == "completed"
|
||||
assert row[1] is None
|
||||
assert row[2] is None
|
||||
@@ -0,0 +1,669 @@
|
||||
"""Tests for the reconciler beat task.
|
||||
|
||||
Seeds stuck rows of each kind (Q1 messages, Q2 proposed tool calls, Q3
|
||||
executed tool calls) and verifies that ``run_reconciliation`` either
|
||||
retries (Q1 within 3 attempts) or escalates them to terminal status
|
||||
with both a structured logger error and a stack_logs row.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers — seed rows in shapes that mirror the live writers (T3 reserve_message
|
||||
# for stuck messages, T8 record_proposed for tool_call_attempts).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _create_conv(conn, user_id: str = "u-1") -> dict:
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
|
||||
return ConversationsRepository(conn).create(user_id, "rec test")
|
||||
|
||||
|
||||
def _seed_pending_message(
|
||||
conn,
|
||||
*,
|
||||
user_id: str = "u-1",
|
||||
age_minutes: int = 6,
|
||||
status: str = "pending",
|
||||
) -> dict:
|
||||
"""Insert a conversation_messages row in ``status`` with stale ``timestamp``."""
|
||||
conv = _create_conv(conn, user_id=user_id)
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO conversation_messages (
|
||||
conversation_id, position, prompt, response, status,
|
||||
user_id, timestamp
|
||||
)
|
||||
VALUES (
|
||||
CAST(:conv_id AS uuid), 0, :prompt, :resp, :status,
|
||||
:user_id,
|
||||
clock_timestamp() - make_interval(mins => :age)
|
||||
)
|
||||
RETURNING id
|
||||
"""
|
||||
),
|
||||
{
|
||||
"conv_id": conv["id"],
|
||||
"prompt": "hello",
|
||||
"resp": "",
|
||||
"status": status,
|
||||
"user_id": user_id,
|
||||
"age": age_minutes,
|
||||
},
|
||||
).fetchone()
|
||||
return {"id": str(row[0]), "conversation_id": conv["id"], "user_id": user_id}
|
||||
|
||||
|
||||
def _seed_resuming_state(conn, conv_id: str, user_id: str, *, secs_ago: int) -> None:
|
||||
"""Insert a pending_tool_state row in ``resuming`` with ``resumed_at`` backdated."""
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO pending_tool_state (
|
||||
conversation_id, user_id, messages, pending_tool_calls,
|
||||
tools_dict, tool_schemas, agent_config,
|
||||
created_at, expires_at, status, resumed_at
|
||||
)
|
||||
VALUES (
|
||||
CAST(:conv_id AS uuid), :user_id,
|
||||
'[]'::jsonb, '[]'::jsonb, '{}'::jsonb, '[]'::jsonb, '{}'::jsonb,
|
||||
clock_timestamp(),
|
||||
clock_timestamp() + interval '30 minutes',
|
||||
'resuming',
|
||||
clock_timestamp() - make_interval(secs => :secs_ago)
|
||||
)
|
||||
"""
|
||||
),
|
||||
{"conv_id": conv_id, "user_id": user_id, "secs_ago": secs_ago},
|
||||
)
|
||||
|
||||
|
||||
def _seed_pending_state(
|
||||
conn, conv_id: str, user_id: str, *, expires_in_minutes: int = 30,
|
||||
) -> None:
|
||||
"""Insert a paused ``pending_tool_state`` row (status='pending').
|
||||
|
||||
A negative ``expires_in_minutes`` simulates a TTL-expired row that
|
||||
the cleanup janitor hasn't yet reaped.
|
||||
"""
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO pending_tool_state (
|
||||
conversation_id, user_id, messages, pending_tool_calls,
|
||||
tools_dict, tool_schemas, agent_config,
|
||||
created_at, expires_at, status, resumed_at
|
||||
)
|
||||
VALUES (
|
||||
CAST(:conv_id AS uuid), :user_id,
|
||||
'[]'::jsonb, '[]'::jsonb, '{}'::jsonb, '[]'::jsonb, '{}'::jsonb,
|
||||
clock_timestamp(),
|
||||
clock_timestamp() + make_interval(mins => :exp),
|
||||
'pending',
|
||||
NULL
|
||||
)
|
||||
"""
|
||||
),
|
||||
{"conv_id": conv_id, "user_id": user_id, "exp": expires_in_minutes},
|
||||
)
|
||||
|
||||
|
||||
def _seed_tool_call(
|
||||
conn,
|
||||
*,
|
||||
call_id: str,
|
||||
status: str,
|
||||
tool_name: str = "notes",
|
||||
age_minutes: int = 16,
|
||||
tool_id: str | None = None,
|
||||
arguments: dict | None = None,
|
||||
action_name: str = "view",
|
||||
) -> None:
|
||||
"""Insert a tool_call_attempts row, then backdate ``attempted_at``/``updated_at``."""
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO tool_call_attempts (
|
||||
call_id, tool_id, tool_name, action_name, arguments, status
|
||||
)
|
||||
VALUES (
|
||||
:call_id, CAST(:tool_id AS uuid), :tool_name, :action_name,
|
||||
CAST(:arguments AS jsonb), :status
|
||||
)
|
||||
"""
|
||||
),
|
||||
{
|
||||
"call_id": call_id,
|
||||
"tool_id": tool_id,
|
||||
"tool_name": tool_name,
|
||||
"action_name": action_name,
|
||||
"arguments": json.dumps(arguments or {}),
|
||||
"status": status,
|
||||
},
|
||||
)
|
||||
# The set_updated_at BEFORE-UPDATE trigger would clobber updated_at;
|
||||
# disable it briefly so the backdate lands.
|
||||
conn.execute(text("ALTER TABLE tool_call_attempts DISABLE TRIGGER USER"))
|
||||
try:
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
UPDATE tool_call_attempts
|
||||
SET attempted_at = clock_timestamp() - make_interval(mins => :age),
|
||||
updated_at = clock_timestamp() - make_interval(mins => :age)
|
||||
WHERE call_id = :call_id
|
||||
"""
|
||||
),
|
||||
{"call_id": call_id, "age": age_minutes},
|
||||
)
|
||||
finally:
|
||||
conn.execute(text("ALTER TABLE tool_call_attempts ENABLE TRIGGER USER"))
|
||||
|
||||
|
||||
def _stack_logs_count(conn, query: str | None = None) -> int:
|
||||
if query is None:
|
||||
result = conn.execute(text("SELECT count(*) FROM stack_logs"))
|
||||
else:
|
||||
result = conn.execute(
|
||||
text("SELECT count(*) FROM stack_logs WHERE query = :q"),
|
||||
{"q": query},
|
||||
)
|
||||
return int(result.scalar() or 0)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _route_engine_to(pg_conn):
|
||||
"""Patch ``get_engine`` so the run_reconciliation reuses the test conn."""
|
||||
|
||||
@contextmanager
|
||||
def _fake_begin():
|
||||
# The repository writes happen on this connection; rolling back
|
||||
# at the end of the test cleans them up.
|
||||
yield pg_conn
|
||||
|
||||
fake_engine = MagicMock()
|
||||
fake_engine.begin = _fake_begin
|
||||
|
||||
with patch(
|
||||
"application.api.user.reconciliation.get_engine",
|
||||
return_value=fake_engine,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Q1 — stuck messages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStuckMessages:
|
||||
@pytest.mark.unit
|
||||
def test_first_two_attempts_increment_only(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r1 = run_reconciliation()
|
||||
r2 = run_reconciliation()
|
||||
|
||||
assert r1["messages_failed"] == 0
|
||||
assert r2["messages_failed"] == 0
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata "
|
||||
"FROM conversation_messages WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
assert row[1]["reconcile_attempts"] == 2
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_third_attempt_marks_failed_and_emits_alert(self, pg_conn, caplog):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
before_logs = _stack_logs_count(pg_conn, "reconciler_message_failed")
|
||||
|
||||
with _route_engine_to(pg_conn), caplog.at_level(
|
||||
logging.ERROR, logger="application.api.user.reconciliation",
|
||||
):
|
||||
run_reconciliation()
|
||||
run_reconciliation()
|
||||
r3 = run_reconciliation()
|
||||
|
||||
assert r3["messages_failed"] == 1
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata "
|
||||
"FROM conversation_messages WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "failed"
|
||||
assert row[1]["reconcile_attempts"] == 3
|
||||
assert "reconciler:" in (row[1].get("error") or "")
|
||||
|
||||
# Structured alert + stack_logs row both surface the failure.
|
||||
assert any(
|
||||
"reconciler alert" in rec.getMessage()
|
||||
and rec.levelname == "ERROR"
|
||||
and getattr(rec, "alert", None) == "reconciler_message_failed"
|
||||
for rec in caplog.records
|
||||
)
|
||||
assert (
|
||||
_stack_logs_count(pg_conn, "reconciler_message_failed")
|
||||
== before_logs + 1
|
||||
)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_streaming_status_also_eligible(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn, status="streaming")
|
||||
with _route_engine_to(pg_conn):
|
||||
run_reconciliation()
|
||||
run_reconciliation()
|
||||
run_reconciliation()
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status FROM conversation_messages "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "failed"
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skipped_when_active_resuming_state(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
# Active resume started 60 seconds ago — within 10-min grace.
|
||||
_seed_resuming_state(pg_conn, msg["conversation_id"], msg["user_id"], secs_ago=60)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["messages_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata "
|
||||
"FROM conversation_messages WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
# No attempts recorded: row was skipped, not just not-yet-failed.
|
||||
assert "reconcile_attempts" not in row[1]
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_stale_resuming_does_not_skip(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
# 11 minutes ago — past the 10-minute grace window.
|
||||
_seed_resuming_state(
|
||||
pg_conn, msg["conversation_id"], msg["user_id"], secs_ago=11 * 60,
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
run_reconciliation()
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT message_metadata FROM conversation_messages "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
# Stale resuming shouldn't shield: attempt counter must have moved.
|
||||
assert row[0].get("reconcile_attempts") == 1
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_fresh_message_left_alone(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
# 1 minute old — well under the 5-minute threshold.
|
||||
msg = _seed_pending_message(pg_conn, age_minutes=1)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["messages_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata FROM conversation_messages "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
assert "reconcile_attempts" not in row[1]
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skipped_when_paused_pending_state_active(self, pg_conn):
|
||||
"""A paused conversation (PT.status='pending') must not get failed
|
||||
while the user is still considering tool approval. The PT row's
|
||||
own ``expires_at`` TTL is the abandonment signal.
|
||||
"""
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
_seed_pending_state(
|
||||
pg_conn, msg["conversation_id"], msg["user_id"],
|
||||
expires_in_minutes=30,
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r1 = run_reconciliation()
|
||||
r2 = run_reconciliation()
|
||||
r3 = run_reconciliation()
|
||||
|
||||
assert r1["messages_failed"] == 0
|
||||
assert r2["messages_failed"] == 0
|
||||
assert r3["messages_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, message_metadata FROM conversation_messages "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
# No attempts recorded: the row was excluded from the sweep, not
|
||||
# just held under the 3-attempt threshold.
|
||||
assert "reconcile_attempts" not in row[1]
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_expired_pending_state_does_not_skip(self, pg_conn):
|
||||
"""When the PT row's TTL has expired but the cleanup janitor
|
||||
hasn't reaped it yet, the message becomes eligible immediately
|
||||
— we don't wait an extra ~60s for janitor cadence to align.
|
||||
"""
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
msg = _seed_pending_message(pg_conn)
|
||||
_seed_pending_state(
|
||||
pg_conn, msg["conversation_id"], msg["user_id"],
|
||||
expires_in_minutes=-1,
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
run_reconciliation()
|
||||
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT message_metadata FROM conversation_messages "
|
||||
"WHERE id = CAST(:id AS uuid)"
|
||||
),
|
||||
{"id": msg["id"]},
|
||||
).fetchone()
|
||||
assert row[0].get("reconcile_attempts") == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Q2 — proposed tool calls (side effect unknown)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStuckProposedToolCalls:
|
||||
@pytest.mark.unit
|
||||
def test_marks_proposed_failed_with_alert(self, pg_conn, caplog):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_tool_call(pg_conn, call_id="cp-1", status="proposed", age_minutes=6)
|
||||
before = _stack_logs_count(pg_conn, "reconciler_tool_call_failed_proposed")
|
||||
|
||||
with _route_engine_to(pg_conn), caplog.at_level(
|
||||
logging.ERROR, logger="application.api.user.reconciliation",
|
||||
):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["tool_calls_failed"] == 1
|
||||
row = pg_conn.execute(
|
||||
text("SELECT status, error FROM tool_call_attempts WHERE call_id = :id"),
|
||||
{"id": "cp-1"},
|
||||
).fetchone()
|
||||
assert row[0] == "failed"
|
||||
assert "side effect status unknown" in (row[1] or "")
|
||||
|
||||
assert any(
|
||||
getattr(rec, "alert", None) == "reconciler_tool_call_failed_proposed"
|
||||
for rec in caplog.records
|
||||
)
|
||||
assert (
|
||||
_stack_logs_count(pg_conn, "reconciler_tool_call_failed_proposed")
|
||||
== before + 1
|
||||
)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_fresh_proposed_left_alone(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_tool_call(pg_conn, call_id="cp-2", status="proposed", age_minutes=2)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["tool_calls_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text("SELECT status FROM tool_call_attempts WHERE call_id = :id"),
|
||||
{"id": "cp-2"},
|
||||
).fetchone()
|
||||
assert row[0] == "proposed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Q3 — executed tool calls (escalate to failed; manual cleanup expected)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStuckExecutedToolCalls:
|
||||
@pytest.mark.unit
|
||||
def test_executed_past_ttl_marked_failed_with_alert(self, pg_conn, caplog):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_tool_call(
|
||||
pg_conn, call_id="ce-1", status="executed",
|
||||
tool_name="notes", age_minutes=16,
|
||||
)
|
||||
before = _stack_logs_count(pg_conn, "reconciler_tool_call_failed_executed")
|
||||
|
||||
with _route_engine_to(pg_conn), caplog.at_level(
|
||||
logging.ERROR, logger="application.api.user.reconciliation",
|
||||
):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["tool_calls_failed"] == 1
|
||||
row = pg_conn.execute(
|
||||
text("SELECT status, error FROM tool_call_attempts WHERE call_id = :id"),
|
||||
{"id": "ce-1"},
|
||||
).fetchone()
|
||||
assert row[0] == "failed"
|
||||
assert "executed-not-confirmed" in (row[1] or "")
|
||||
assert "manual cleanup" in (row[1] or "")
|
||||
assert any(
|
||||
getattr(rec, "alert", None) == "reconciler_tool_call_failed_executed"
|
||||
for rec in caplog.records
|
||||
)
|
||||
assert (
|
||||
_stack_logs_count(pg_conn, "reconciler_tool_call_failed_executed")
|
||||
== before + 1
|
||||
)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_fresh_executed_left_alone(self, pg_conn):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_tool_call(
|
||||
pg_conn, call_id="ce-2", status="executed",
|
||||
tool_name="notes", age_minutes=5, # under 15-min threshold
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["tool_calls_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text("SELECT status FROM tool_call_attempts WHERE call_id = :id"),
|
||||
{"id": "ce-2"},
|
||||
).fetchone()
|
||||
assert row[0] == "executed"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Q5 — stuck idempotency pending rows (lease expired + attempts exhausted)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _seed_stuck_idempotency_row(
|
||||
conn,
|
||||
*,
|
||||
key: str,
|
||||
attempt_count: int,
|
||||
lease_secs_ago: int,
|
||||
) -> None:
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO task_dedup (
|
||||
idempotency_key, task_name, task_id, status,
|
||||
attempt_count, lease_owner_id, lease_expires_at,
|
||||
created_at
|
||||
) VALUES (
|
||||
:key, 'ingest', :tid, 'pending', :attempts,
|
||||
:owner,
|
||||
clock_timestamp() - make_interval(secs => :secs),
|
||||
clock_timestamp() - make_interval(mins => 5)
|
||||
)
|
||||
"""
|
||||
),
|
||||
{
|
||||
"key": key,
|
||||
"tid": f"task-{key}",
|
||||
"attempts": int(attempt_count),
|
||||
"owner": f"owner-{key}",
|
||||
"secs": int(lease_secs_ago),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TestStuckIdempotencyPending:
|
||||
@pytest.mark.unit
|
||||
def test_promotes_to_failed_with_alert(self, pg_conn, caplog):
|
||||
"""A pending row whose lease expired and whose attempt counter
|
||||
already hit the poison-loop threshold gets escalated to failed
|
||||
so a same-key retry can re-claim instead of waiting 24 h.
|
||||
"""
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_stuck_idempotency_row(
|
||||
pg_conn, key="abandoned", attempt_count=5, lease_secs_ago=120,
|
||||
)
|
||||
before = _stack_logs_count(
|
||||
pg_conn, "reconciler_idempotency_pending_failed",
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn), caplog.at_level(
|
||||
logging.ERROR, logger="application.api.user.reconciliation",
|
||||
):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["idempotency_pending_failed"] == 1
|
||||
row = pg_conn.execute(
|
||||
text(
|
||||
"SELECT status, result_json FROM task_dedup "
|
||||
"WHERE idempotency_key = :k"
|
||||
),
|
||||
{"k": "abandoned"},
|
||||
).fetchone()
|
||||
assert row[0] == "failed"
|
||||
assert row[1]["reconciled"] is True
|
||||
assert "lease expired" in row[1]["error"]
|
||||
|
||||
assert any(
|
||||
getattr(rec, "alert", None) == "reconciler_idempotency_pending_failed"
|
||||
for rec in caplog.records
|
||||
)
|
||||
assert (
|
||||
_stack_logs_count(
|
||||
pg_conn, "reconciler_idempotency_pending_failed",
|
||||
)
|
||||
== before + 1
|
||||
)
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skips_under_attempt_threshold(self, pg_conn):
|
||||
"""Attempt count below the threshold means the wrapper might
|
||||
still re-claim cleanly — leave the row alone.
|
||||
"""
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_stuck_idempotency_row(
|
||||
pg_conn, key="recoverable", attempt_count=2, lease_secs_ago=120,
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["idempotency_pending_failed"] == 0
|
||||
row = pg_conn.execute(
|
||||
text("SELECT status FROM task_dedup WHERE idempotency_key = :k"),
|
||||
{"k": "recoverable"},
|
||||
).fetchone()
|
||||
assert row[0] == "pending"
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skips_within_lease_grace(self, pg_conn):
|
||||
"""A lease that just expired (10 s ago) might be in the
|
||||
heartbeat-tick window; the 60 s grace keeps the sweep quiet.
|
||||
"""
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
|
||||
_seed_stuck_idempotency_row(
|
||||
pg_conn, key="just-expired", attempt_count=5, lease_secs_ago=10,
|
||||
)
|
||||
|
||||
with _route_engine_to(pg_conn):
|
||||
r = run_reconciliation()
|
||||
|
||||
assert r["idempotency_pending_failed"] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Skip path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPostgresUriMissing:
|
||||
@pytest.mark.unit
|
||||
def test_returns_skip_dict(self, monkeypatch):
|
||||
from application.api.user.reconciliation import run_reconciliation
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False)
|
||||
|
||||
result = run_reconciliation()
|
||||
assert result == {
|
||||
"messages_failed": 0,
|
||||
"tool_calls_failed": 0,
|
||||
"skipped": "POSTGRES_URI not set",
|
||||
}
|
||||
@@ -1,9 +1,26 @@
|
||||
from contextlib import contextmanager
|
||||
from datetime import timedelta
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patch_decorator_db(conn):
|
||||
"""Route the decorator's own ``db_session`` / ``db_readonly`` at ``conn``."""
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch(
|
||||
"application.api.user.idempotency.db_session", _yield
|
||||
), patch(
|
||||
"application.api.user.idempotency.db_readonly", _yield
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
class TestIngestTask:
|
||||
@pytest.mark.unit
|
||||
@patch("application.api.user.tasks.ingest_worker")
|
||||
@@ -16,7 +33,7 @@ class TestIngestTask:
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
|
||||
file_name_map=None,
|
||||
file_name_map=None, idempotency_key=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -33,7 +50,7 @@ class TestIngestTask:
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, "dir", ["pdf"], "job1", "/path", "file.pdf", "user1",
|
||||
file_name_map=name_map,
|
||||
file_name_map=name_map, idempotency_key=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -48,7 +65,8 @@ class TestIngestRemoteTask:
|
||||
result = ingest_remote({"url": "http://x"}, "job1", "user1", "web")
|
||||
|
||||
mock_worker.assert_called_once_with(
|
||||
ANY, {"url": "http://x"}, "job1", "user1", "web"
|
||||
ANY, {"url": "http://x"}, "job1", "user1", "web",
|
||||
idempotency_key=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -150,6 +168,7 @@ class TestIngestConnectorTask:
|
||||
operation_mode="upload",
|
||||
doc_id=None,
|
||||
sync_frequency="never",
|
||||
idempotency_key=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -187,6 +206,7 @@ class TestIngestConnectorTask:
|
||||
operation_mode="sync",
|
||||
doc_id="doc1",
|
||||
sync_frequency="daily",
|
||||
idempotency_key=None,
|
||||
)
|
||||
assert result == {"status": "ok"}
|
||||
|
||||
@@ -200,7 +220,7 @@ class TestSetupPeriodicTasks:
|
||||
|
||||
setup_periodic_tasks(sender)
|
||||
|
||||
assert sender.add_periodic_task.call_count == 5
|
||||
assert sender.add_periodic_task.call_count == 7
|
||||
|
||||
calls = sender.add_periodic_task.call_args_list
|
||||
|
||||
@@ -212,8 +232,15 @@ class TestSetupPeriodicTasks:
|
||||
assert calls[2][0][0] == timedelta(days=30)
|
||||
# pending_tool_state TTL cleanup (60s)
|
||||
assert calls[3][0][0] == timedelta(seconds=60)
|
||||
assert calls[3][1].get("name") == "cleanup-pending-tool-state"
|
||||
# idempotency dedup TTL cleanup (1h)
|
||||
assert calls[4][0][0] == timedelta(hours=1)
|
||||
assert calls[4][1].get("name") == "cleanup-idempotency-dedup"
|
||||
# reconciliation sweep (30s)
|
||||
assert calls[5][0][0] == timedelta(seconds=30)
|
||||
assert calls[5][1].get("name") == "reconciliation"
|
||||
# version-check (every 7h)
|
||||
assert calls[4][0][0] == timedelta(hours=7)
|
||||
assert calls[6][0][0] == timedelta(hours=7)
|
||||
|
||||
|
||||
class TestMcpOauthTask:
|
||||
@@ -242,3 +269,203 @@ class TestMcpOauthStatusTask:
|
||||
|
||||
mock_worker.assert_called_once_with(ANY, "task123")
|
||||
assert result == {"status": "authorized"}
|
||||
|
||||
|
||||
class TestDurableTaskRetryPolicy:
|
||||
"""The long-running tasks share a uniform retry policy."""
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.parametrize(
|
||||
"task_name",
|
||||
[
|
||||
"ingest",
|
||||
"ingest_remote",
|
||||
"reingest_source_task",
|
||||
"store_attachment",
|
||||
"process_agent_webhook",
|
||||
"ingest_connector_task",
|
||||
],
|
||||
)
|
||||
def test_task_has_retry_config(self, task_name):
|
||||
import application.api.user.tasks as tasks_module
|
||||
|
||||
task = getattr(tasks_module, task_name)
|
||||
assert task.acks_late is True
|
||||
assert Exception in task.autoretry_for
|
||||
assert task.retry_backoff is True
|
||||
assert task.retry_kwargs == {"max_retries": 3, "countdown": 60}
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.parametrize(
|
||||
"task_name",
|
||||
[
|
||||
"schedule_syncs",
|
||||
"sync_source",
|
||||
"mcp_oauth_task",
|
||||
"mcp_oauth_status_task",
|
||||
"cleanup_pending_tool_state",
|
||||
"reconciliation_task",
|
||||
"version_check_task",
|
||||
],
|
||||
)
|
||||
def test_short_periodic_tasks_have_no_retry_config(self, task_name):
|
||||
import application.api.user.tasks as tasks_module
|
||||
|
||||
task = getattr(tasks_module, task_name)
|
||||
assert not getattr(task, "autoretry_for", None)
|
||||
|
||||
|
||||
class TestProcessAgentWebhookIdempotency:
|
||||
"""Wrapper short-circuits a second call with the same key on the durable webhook task."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_repeat_with_same_key_short_circuits(self, pg_conn):
|
||||
from application.api.user.tasks import process_agent_webhook
|
||||
|
||||
worker_calls = []
|
||||
|
||||
def _fake_worker(self, agent_id, payload):
|
||||
worker_calls.append((agent_id, payload))
|
||||
return {"status": "success", "result": {"answer": "ok"}}
|
||||
|
||||
with _patch_decorator_db(pg_conn), patch(
|
||||
"application.api.user.tasks.agent_webhook_worker",
|
||||
side_effect=_fake_worker,
|
||||
):
|
||||
first = process_agent_webhook(
|
||||
"agent", {"event": "x"}, idempotency_key="dur-k1",
|
||||
)
|
||||
second = process_agent_webhook(
|
||||
"agent", {"event": "x"}, idempotency_key="dur-k1",
|
||||
)
|
||||
|
||||
assert first == {"status": "success", "result": {"answer": "ok"}}
|
||||
assert second == first
|
||||
assert len(worker_calls) == 1
|
||||
|
||||
|
||||
class TestCleanupPendingToolState:
|
||||
"""Janitor reverts stale 'resuming' rows and deletes TTL-expired rows."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_reverts_stale_and_deletes_expired(self, pg_conn):
|
||||
from sqlalchemy import text as _text
|
||||
|
||||
from application.api.user.tasks import cleanup_pending_tool_state
|
||||
from application.storage.db.repositories.conversations import (
|
||||
ConversationsRepository,
|
||||
)
|
||||
from application.storage.db.repositories.pending_tool_state import (
|
||||
PendingToolStateRepository,
|
||||
)
|
||||
|
||||
repo = PendingToolStateRepository(pg_conn)
|
||||
|
||||
def _sample() -> dict:
|
||||
return {
|
||||
"messages": [],
|
||||
"pending_tool_calls": [],
|
||||
"tools_dict": {},
|
||||
"tool_schemas": [],
|
||||
"agent_config": {},
|
||||
}
|
||||
|
||||
# Pending and fresh — should be left alone.
|
||||
c1 = ConversationsRepository(pg_conn).create("u", "fresh-pending")
|
||||
repo.save_state(c1["id"], "u", **_sample())
|
||||
|
||||
# Pending but already expired — should be deleted.
|
||||
c2 = ConversationsRepository(pg_conn).create("u", "expired-pending")
|
||||
repo.save_state(c2["id"], "u", **_sample(), ttl_seconds=0)
|
||||
|
||||
# Resuming within grace — should stay 'resuming'.
|
||||
c3 = ConversationsRepository(pg_conn).create("u", "fresh-resuming")
|
||||
repo.save_state(c3["id"], "u", **_sample())
|
||||
repo.mark_resuming(c3["id"], "u")
|
||||
|
||||
# Resuming past grace — should revert to 'pending'.
|
||||
c4 = ConversationsRepository(pg_conn).create("u", "stale-resuming")
|
||||
repo.save_state(c4["id"], "u", **_sample())
|
||||
repo.mark_resuming(c4["id"], "u")
|
||||
pg_conn.execute(
|
||||
_text(
|
||||
"UPDATE pending_tool_state "
|
||||
"SET resumed_at = clock_timestamp() "
|
||||
" - make_interval(secs => 660) "
|
||||
"WHERE conversation_id = CAST(:conv_id AS uuid)"
|
||||
),
|
||||
{"conv_id": c4["id"]},
|
||||
)
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def _fake_begin():
|
||||
yield pg_conn
|
||||
|
||||
fake_engine = MagicMock()
|
||||
fake_engine.begin = _fake_begin
|
||||
|
||||
with patch(
|
||||
"application.storage.db.engine.get_engine",
|
||||
return_value=fake_engine,
|
||||
):
|
||||
result = cleanup_pending_tool_state.run()
|
||||
|
||||
assert result["reverted"] == 1
|
||||
assert result["deleted"] == 1
|
||||
|
||||
# Final state assertions.
|
||||
assert repo.load_state(c1["id"], "u")["status"] == "pending"
|
||||
assert repo.load_state(c2["id"], "u") is None
|
||||
assert repo.load_state(c3["id"], "u")["status"] == "resuming"
|
||||
c4_row = repo.load_state(c4["id"], "u")
|
||||
assert c4_row["status"] == "pending"
|
||||
assert c4_row["resumed_at"] is None
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_skips_when_postgres_uri_missing(self, monkeypatch):
|
||||
from application.api.user.tasks import cleanup_pending_tool_state
|
||||
from application.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "POSTGRES_URI", None, raising=False)
|
||||
|
||||
result = cleanup_pending_tool_state.run()
|
||||
assert result == {
|
||||
"deleted": 0,
|
||||
"reverted": 0,
|
||||
"skipped": "POSTGRES_URI not set",
|
||||
}
|
||||
|
||||
|
||||
class TestIngestIdempotency:
|
||||
"""Same short-circuit applies to the ingest task path."""
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_repeat_with_same_key_short_circuits(self, pg_conn):
|
||||
from application.api.user.tasks import ingest
|
||||
|
||||
worker_calls = []
|
||||
|
||||
def _fake_worker(self, directory, formats, job_name, file_path,
|
||||
filename, user, file_name_map=None,
|
||||
idempotency_key=None):
|
||||
worker_calls.append(filename)
|
||||
return {"status": "ok", "directory": directory}
|
||||
|
||||
with _patch_decorator_db(pg_conn), patch(
|
||||
"application.api.user.tasks.ingest_worker",
|
||||
side_effect=_fake_worker,
|
||||
):
|
||||
first = ingest(
|
||||
"dir", ["pdf"], "job1", "user1", "/path", "file.pdf",
|
||||
idempotency_key="dur-ing-1",
|
||||
)
|
||||
second = ingest(
|
||||
"dir", ["pdf"], "job1", "user1", "/path", "file.pdf",
|
||||
idempotency_key="dur-ing-1",
|
||||
)
|
||||
|
||||
assert first == second
|
||||
assert first == {"status": "ok", "directory": "dir"}
|
||||
assert len(worker_calls) == 1
|
||||
Reference in new issue
Block a user