diff --git a/tests/integration/test_sources_jobs.py b/tests/integration/test_sources_jobs.py index 5d761f2..69d7d5e 100644 --- a/tests/integration/test_sources_jobs.py +++ b/tests/integration/test_sources_jobs.py @@ -1,5 +1,6 @@ from __future__ import annotations +from datetime import UTC, datetime, timedelta from pathlib import Path import httpx @@ -7,6 +8,13 @@ import pytest from backup_tool.api.app import create_app from backup_tool.config import Settings from backup_tool.db.models import Execution +from backup_tool.execution import ( + claim, + complete_cancellation, + heartbeat, + recover_stale, + request_cancellation, +) PASSWORD = "correct-horse-battery-staple" @@ -95,14 +103,24 @@ async def test_local_source_probe_archive_and_repository_targeted_job( ) assert duplicate.status_code == 409 assert duplicate.json()["code"] == "execution_active" - polled = await client.get(f"/api/v2/executions/{execution.json()['id']}", headers=headers) + execution_id = execution.json()["id"] + async with app.state.sessions() as db: + assert await claim(db, execution_id, "expired-worker") is not None + stored = await db.get(Execution, execution_id) + assert stored is not None + stored.lease_expires_at = datetime.now(UTC) - timedelta(seconds=1) + await db.commit() + async with app.state.sessions() as db: + assert await recover_stale(db) == 1 + assert await claim(db, execution_id, "replacement-worker") is not None + async with app.state.sessions() as db: + assert not await heartbeat(db, execution_id, "expired-worker") + assert await request_cancellation(db, execution_id) is not None + assert not await complete_cancellation(db, execution_id, "expired-worker") + assert await complete_cancellation(db, execution_id, "replacement-worker") + polled = await client.get(f"/api/v2/executions/{execution_id}", headers=headers) assert polled.status_code == 200 - assert polled.json()["state"] == "queued" - cancellation = await client.post( - f"/api/v2/executions/{execution.json()['id']}/cancel", headers=headers - ) - assert cancellation.status_code == 202 - assert cancellation.json()["state"] == "cancelled" + assert polled.json()["state"] == "cancelled" async with app.state.sessions() as db: stored = await db.get(Execution, execution.json()["id"]) assert stored is not None