142 lines
6.6 KiB
Python
142 lines
6.6 KiB
Python
import json
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from scripts.run_operational_worker import CredentialProvider, OperationalWorker, WorkerSettings
|
|
|
|
|
|
@pytest.fixture
|
|
def anyio_backend():
|
|
return "asyncio"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_persisted_queue_retries_with_backend_execution_identity():
|
|
task_id, run_id, execution_id = str(uuid4()), str(uuid4()), str(uuid4())
|
|
submitted = []
|
|
def handler(request):
|
|
assert request.headers["Authorization"] == "Bearer TEST_ONLY"
|
|
assert not request.url.path.endswith("schedules/dispatch-due")
|
|
if request.url.path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
if request.method == "GET":
|
|
if request.url.params.get("status") != "queued":
|
|
return httpx.Response(200, json={"items": []})
|
|
return httpx.Response(200, json={"items": [{
|
|
"task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
|
|
"status": "queued", "revision": 1,
|
|
}]})
|
|
body = json.loads(request.content)
|
|
submitted.append(body)
|
|
return httpx.Response(202, json={"task_id": task_id, "run_id": run_id, "session_id": "execution-session"})
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
worker = OperationalWorker(settings, client)
|
|
assert await worker.tick() == 1
|
|
assert await worker.tick() == 1
|
|
assert submitted[0] == submitted[1]
|
|
assert submitted[0] == {
|
|
"task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
|
|
}
|
|
assert "credential" not in submitted[0]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("run_status", ["claimed", "running", "failed"])
|
|
async def test_only_queued_runs_are_submitted(run_status):
|
|
def handler(request):
|
|
assert not request.url.path.endswith("schedules/dispatch-due")
|
|
if request.method == "GET":
|
|
requested = request.url.params.get("status")
|
|
if requested == run_status and run_status in {"claimed", "running"}:
|
|
return httpx.Response(200, json={"items": [{"run_id": str(uuid4()), "status": run_status}]})
|
|
return httpx.Response(200, json={"items": []})
|
|
assert request.url.host == "server.test"
|
|
return httpx.Response(200, json={"items": []})
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 0
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_bounded_queue_continues_after_one_agent_handoff_fails():
|
|
runs = [{"task_id": str(uuid4()), "run_id": str(uuid4()),
|
|
"execution_request_id": str(uuid4()), "status": "queued",
|
|
"revision": 1} for _ in range(2)]
|
|
limits, submitted, failures = [], [], []
|
|
|
|
def handler(request):
|
|
if request.method == "GET":
|
|
if request.url.params.get("status") != "queued":
|
|
return httpx.Response(200, json={"items": []})
|
|
limits.append(int(request.url.params["limit"]))
|
|
return httpx.Response(200, json={"items": runs})
|
|
if request.url.host == "server.test":
|
|
if request.url.path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
failures.append(json.loads(request.content))
|
|
return httpx.Response(200, json={})
|
|
body = json.loads(request.content)
|
|
submitted.append(body)
|
|
if body["task_id"] == runs[0]["task_id"]:
|
|
return httpx.Response(503, json={"detail": "temporary failure"})
|
|
return httpx.Response(202, json={"task_id": body["task_id"], "run_id": body["run_id"]})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 1
|
|
assert limits == [settings.max_concurrency]
|
|
assert [body["task_id"] for body in submitted] == [run["task_id"] for run in runs]
|
|
assert [body["execution_request_id"] for body in submitted] == [run["execution_request_id"] for run in runs]
|
|
assert failures[0]["retryable"] is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_service_credentials_are_cached_and_refreshed_without_entering_task_payload(monkeypatch):
|
|
clock = [100.0]
|
|
monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
|
|
requests = []
|
|
|
|
def handler(request):
|
|
requests.append(request)
|
|
assert request.url == "https://identity.test/token"
|
|
assert "authorization" not in request.headers
|
|
return httpx.Response(200, json={"access_token": f"TEST_ONLY_{len(requests)}", "expires_in": 60})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
token_url="https://identity.test/token", client_id="test-service",
|
|
client_secret="TEST_ONLY_SECRET")
|
|
provider = CredentialProvider(settings)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
first = await provider.headers(client)
|
|
clock[0] = 140.0
|
|
assert await provider.headers(client) == first
|
|
assert len(requests) == 1
|
|
clock[0] = 146.0
|
|
refreshed = await provider.headers(client)
|
|
assert len(requests) == 2
|
|
assert first["Authorization"] == "Bearer TEST_ONLY_1"
|
|
assert refreshed == {"Authorization": "Bearer TEST_ONLY_2", "X-Project-Id": settings.project_id}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_long_lived_oauth_token_is_refreshed_before_product_token_age_limit(monkeypatch):
|
|
clock, requests = [0.0], []
|
|
monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
|
|
|
|
def handler(request):
|
|
requests.append(request)
|
|
return httpx.Response(200, json={"access_token": "TEST_ONLY", "expires_in": 3600})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
token_url="https://identity.test/token", client_id="test-service",
|
|
client_secret="TEST_ONLY_SECRET")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
provider = CredentialProvider(settings)
|
|
await provider.headers(client)
|
|
clock[0] = 301.0
|
|
await provider.headers(client)
|
|
assert len(requests) == 2
|