diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py new file mode 100644 index 0000000..2ad180b --- /dev/null +++ b/tests/lib/qpu_client/test_auth.py @@ -0,0 +1,247 @@ +"""Testing lib/qpu_client/auth""" + +import json +import logging + +import httpx +import pytest +from pytest_httpx import HTTPXMock + +from warden.lib.config.config import QPUAuthConfig, QPUConfig +from warden.lib.qpu_client.auth import ( + KeycloakClientCredentialsAuth, + TokenRequestError, +) + +TOKEN_URL = "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" +QPU_URL = "http://qpu:4300/api/v1/system" + + +@pytest.fixture +def auth_conf() -> QPUAuthConfig: + return QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s3cret" + ) + + +def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, + json={"access_token": "tok-1", "expires_in": 300}, + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + first = client.get(QPU_URL) + second = client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-1" + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + + +def test_token_request_uses_client_credentials_grant(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + client.get(QPU_URL) + + token_request = next( + r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL + ) + body = token_request.read().decode() + assert "grant_type=client_credentials" in body + assert "client_id=warden" in body + assert "client_secret=s3cret" in body + + +def test_expired_token_is_refreshed(httpx_mock: HTTPXMock, auth_conf, monkeypatch): + # expires_in 300 with leeway 30 means the token is stale after 270s. + clock = {"now": 1_000.0} + monkeypatch.setattr("warden.lib.qpu_client.auth.monotonic", lambda: clock["now"]) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + first = client.get(QPU_URL) + clock["now"] += 271 + second = client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-2" + + +def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + response = client.get(QPU_URL) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + response = client.get(QPU_URL) + + # The second 401 is surfaced, not retried again. + assert response.status_code == 401 + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +@pytest.mark.parametrize("status_code", [400, 401]) +def test_bad_credentials_raise_token_request_error( + httpx_mock: HTTPXMock, auth_conf, status_code +): + httpx_mock.add_response( + url=TOKEN_URL, + status_code=status_code, + json={"error": "invalid_client"}, + ) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + with pytest.raises(TokenRequestError, match="invalid_client"): + client.get(QPU_URL) + + +def test_keycloak_5xx_raises_retryable_http_status_error( + httpx_mock: HTTPXMock, auth_conf +): + # 503 must stay an httpx.HTTPStatusError so the existing retry decorator + # recognises it as transient. + httpx_mock.add_response(url=TOKEN_URL, status_code=503) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + with pytest.raises(httpx.HTTPStatusError): + client.get(QPU_URL) + + +@pytest.mark.asyncio +async def test_async_flow_attaches_token(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-async", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with httpx.AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.request.headers["Authorization"] == "Bearer tok-async" + + +@pytest.mark.asyncio +async def test_async_flow_reuses_token_cached_by_sync_flow( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "shared", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + client.get(QPU_URL) + async with httpx.AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.request.headers["Authorization"] == "Bearer shared" + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + + +def test_short_lived_token_still_caches_with_warning( + httpx_mock: HTTPXMock, auth_conf, caplog +): + # expires_in 30 with the default leeway_s 30 would clamp to 0 without the + # half-lifespan fallback, disabling the cache entirely and forcing a + # Keycloak round-trip on every request. + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 30} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with caplog.at_level(logging.WARNING, logger="warden.lib.qpu_client.auth"): + with httpx.Client(auth=auth) as client: + client.get(QPU_URL) + client.get(QPU_URL) + + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + assert any(record.levelno == logging.WARNING for record in caplog.records) + + +def test_client_sends_no_authorization_header_without_auth_config( + httpx_mock: HTTPXMock, +): + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + response = QPUConfig(uri="http://qpu:4300").client.get(QPU_URL) + + assert "Authorization" not in response.request.headers + + +def test_401_on_post_retries_with_fresh_token_and_identical_body( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + body = {"circuit": "bell", "shots": 100} + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + response = client.post(QPU_URL, json=body) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + for request in qpu_requests: + assert json.loads(request.read()) == body diff --git a/tests/lib/qpu_client/test_retry.py b/tests/lib/qpu_client/test_retry.py new file mode 100644 index 0000000..a9ad81c --- /dev/null +++ b/tests/lib/qpu_client/test_retry.py @@ -0,0 +1,30 @@ +"""Testing lib/qpu_client/retry""" + +import pytest + +from warden.lib.qpu_client.auth import TokenRequestError +from warden.lib.qpu_client.retry import UnhandledError, retry + + +def test_already_classified_errors_are_not_rewrapped(): + calls = {"n": 0} + + @retry(max=5, sleep_s=0) + def fails_with_bad_credentials(): + calls["n"] += 1 + raise TokenRequestError("invalid_client") + + with pytest.raises(TokenRequestError): + fails_with_bad_credentials() + + # Fail fast: a wrong secret will not fix itself. + assert calls["n"] == 1 + + +def test_unknown_errors_are_still_wrapped(): + @retry(max=5, sleep_s=0) + def fails_with_value_error(): + raise ValueError("something unexpected") + + with pytest.raises(UnhandledError): + fails_with_value_error() diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 697934c..22d415b 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -18,7 +18,8 @@ from warden.lib.config import Config, SchedulerStrategy from warden.lib.models import Job from warden.scheduler.main import run_scheduler -from warden.scheduler.worker import LocalQPUWorker +from warden.scheduler.types import JobUpdateQueue +from warden.scheduler.worker import TERMINAL_STATUSES, LocalQPUWorker NOW = datetime.now() @@ -35,8 +36,6 @@ SYSTEM_API = API_URI + "/system" PROGRAM_API = API_URI + "/programs" -SUCCESS_CHECK_INTERVAL_S = 0.1 - DUMMY_RESULTS = json.dumps([{"counter": {"0001": 1, "0010": 2, "0100": 3, "1000": 4}}]) @@ -60,7 +59,7 @@ async def test_run_nominal( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS - Check "DONE" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -73,7 +72,6 @@ async def test_run_nominal( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 conf: Config = build_conf(strategy, QPU_URI) @@ -147,8 +145,6 @@ async def test_run_nominal( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -157,13 +153,7 @@ async def test_run_nominal( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -200,7 +190,7 @@ async def test_run_resume_job( - To return "DONE" status for ALREADY_DONE_BACKEND_ID - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = 3 - Check "DONE" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -213,7 +203,6 @@ async def test_run_resume_job( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 5 NORMAL_BACKEND_ID = "1" NON_EXISTING_BACKEND_ID = "9999" NEW_BACKEND_ID = "2" @@ -330,8 +319,6 @@ async def test_run_resume_job( ], ) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -340,13 +327,7 @@ async def test_run_resume_job( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -377,7 +358,7 @@ async def test_run_qpu_down( - No need to mock jobs calls - Run scheduler until: - All jobs have an "ERROR" status - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS - Check those jobs have non-empty logs """ @@ -389,7 +370,6 @@ async def test_run_qpu_down( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -413,8 +393,6 @@ async def test_run_qpu_down( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) - ################## ### TEST RUN ### ################## @@ -423,13 +401,9 @@ async def test_run_qpu_down( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_STATUS) all_jobs = (await session.execute(stmt_all)).scalars().all() @@ -495,7 +469,6 @@ async def test_run_job_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 8 N_JOBS_TIMEOUT = 4 @@ -648,13 +621,9 @@ async def test_run_job_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_processed, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) n_processed = (await session.execute(stmt_processed)).scalar() assert n_processed == N_JOBS @@ -688,7 +657,7 @@ async def test_run_resume_job_timeout( - Accept the job's cancelation request - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "CANCELED") == 1 - Check "CANCELED" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -701,7 +670,6 @@ async def test_run_resume_job_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 5 N_JOBS = 1 BACKEND_ID = "1" # Setting the job's created_at at a time that is already timedout @@ -771,8 +739,6 @@ async def test_run_resume_job_timeout( backend_ids=[BACKEND_ID], ) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_JOB_STATUS) - ################## ### TEST RUN ### ################## @@ -781,13 +747,9 @@ async def test_run_resume_job_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_JOB_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_JOB_STATUS) jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -823,7 +785,7 @@ async def test_run_retry_transient_errors( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks: - n (jobs with status "DONE") = N_JOBS - jobs have non-empty logs @@ -836,7 +798,6 @@ async def test_run_retry_transient_errors( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 conf: Config = build_conf(strategy, QPU_URI) @@ -929,7 +890,6 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") stmt = select(Job).where(Job.status == "DONE") ################## @@ -940,13 +900,7 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) jobs_done = (await session.execute(stmt)).scalars().all() assert len(jobs_done) == N_JOBS @@ -974,7 +928,7 @@ async def test_run_qpu_api_unreachable( - To return QPU status as "Down" - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -987,7 +941,6 @@ async def test_run_qpu_api_unreachable( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1003,7 +956,6 @@ async def test_run_qpu_api_unreachable( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1014,13 +966,9 @@ async def test_run_qpu_api_unreachable( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) error_jobs = (await session.execute(stmt)).scalars().all() assert len(error_jobs) == N_JOBS @@ -1050,7 +998,7 @@ async def test_run_job_creation_client_error( - Return exceptions when attempting to create a job - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -1063,7 +1011,6 @@ async def test_run_job_creation_client_error( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -1092,7 +1039,6 @@ async def test_run_job_creation_client_error( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1103,13 +1049,9 @@ async def test_run_job_creation_client_error( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS @@ -1144,7 +1086,7 @@ async def test_run_job_client_error_timeout( (it's the same backend request in QPU) - Run scheduler until: - All jobs have an "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS """ @@ -1155,7 +1097,6 @@ async def test_run_job_client_error_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1208,7 +1149,6 @@ async def test_run_job_client_error_timeout( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1219,13 +1159,9 @@ async def test_run_job_client_error_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS @@ -1256,7 +1192,7 @@ async def test_run_job_canceled_by_cancellation_worker( - For JOB_ID_CANCELED return "CANCELED" status - Run scheduler until: - All jobs have a "DONE" or "CANCELED" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS-1 - Check "DONE" jobs have the right results and non-empty logs - Check n (jobs with status "CANCELED") = 1 @@ -1270,7 +1206,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 JOB_ID_CANCELED = 5 @@ -1376,8 +1311,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status.in_(("DONE", "CANCELED"))) - ################## ### TEST RUN ### ################## @@ -1386,13 +1319,9 @@ async def test_run_job_canceled_by_cancellation_worker( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) stmt_done = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_done)).scalars().all() @@ -1547,3 +1476,61 @@ async def crash(*args, **kwargs): job = (await session.execute(select(Job))).scalar_one() assert job.backend_id == "0" assert job.status == "RUNNING" + + +@pytest.mark.asyncio +async def test_terminal_status_and_closing_log_are_committed_together( + httpx_mock: HTTPXMock, + caplog, +): + """A job's terminal status must be queued together with its closing log line. + + One JobUpdate is one transaction, so pushing the status and the closing + "Job execution ended with status '...'" line separately leaves a window + where the DB holds a finished job whose logs are truncated. Anything that + stops polling once the status is terminal - the API, and every test that + waits on a status then asserts on logs - reads incomplete logs. + + This guards the invariant: keep the closing line in the same update as the + terminal status. + """ + + # Enable warden logging for jobs 'logs' field to be populated + caplog.set_level(logging.INFO, logger="warden") + + def job_json(status: str) -> dict: + return { + "data": { + "uid": 0, + "batch_id": SLURM_USER_ID, + "status": status, + "result": DUMMY_RESULTS if status == "DONE" else None, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": (NOW + timedelta(seconds=1)).isoformat(), + "end_datetime": (NOW + timedelta(seconds=2)).isoformat(), + } + } + + httpx_mock.add_response( + method="GET", + url=SYSTEM_OPERATIONAL_API, + json={"data": {"operational_status": "UP"}}, + ) + httpx_mock.add_response( + method="POST", status_code=200, url=JOB_API, json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("DONE") + ) + + queue: JobUpdateQueue = JobUpdateQueue() + worker = LocalQPUWorker(conf=build_conf(SchedulerStrategy.FIFO, QPU_URI)) + await worker.execute_job(queue=queue, nb_run=100, sequence="{}") + + updates = [queue.get_nowait() for _ in range(queue.qsize())] + first_terminal = next(u for u in updates if u.status in TERMINAL_STATUSES) + assert "Job execution ended with status 'DONE'" in first_terminal.new_logs diff --git a/tests/scheduler/test_scheduler_integration.py b/tests/scheduler/test_scheduler_integration.py index dc9f861..4b6bb36 100644 --- a/tests/scheduler/test_scheduler_integration.py +++ b/tests/scheduler/test_scheduler_integration.py @@ -76,12 +76,10 @@ async def test_run_scheduler_integration( await utils.create_n_jobs(db_session_maker, N_JOBS) - # The terminal status is committed before the closing "Job execution ended - # with status 'DONE'" log line is flushed, so waiting on the status alone - # races with the log assertions below. Wait for the logs too. - stmt = select(func.count(Job.id)).where( - Job.status == "DONE", Job.logs.contains("DONE") - ) + # Safe to wait on the status alone: the scheduler commits a job's terminal + # status and its complete logs in the same transaction, so the log + # assertions below cannot race it + stmt = select(func.count(Job.id)).where(Job.status == "DONE") ################## ### TEST RUN ### @@ -187,11 +185,8 @@ async def test_run_scheduler_integration_cancellation_worker( JOB_TO_CANCEL_ID = job_to_cancel.id - # Same race as above: wait for the closing log line, not just the status - stmt = select(func.count(Job.id)).where( - Job.status.in_(("CANCELED", "DONE")), - Job.logs.contains("Job execution ended"), - ) + # Status alone is enough here too, see the comment in the test above + stmt = select(func.count(Job.id)).where(Job.status.in_(("CANCELED", "DONE"))) ################## ### TEST RUN ### diff --git a/tests/scheduler/utils.py b/tests/scheduler/utils.py index 255503d..a925c91 100644 --- a/tests/scheduler/utils.py +++ b/tests/scheduler/utils.py @@ -2,15 +2,22 @@ import asyncio from asyncio import Task, timeout +from collections.abc import Sequence from contextlib import asynccontextmanager from typing import Any import pytest +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from warden.lib.config import Config, QPUConfig, SchedulerConfig, SchedulerStrategy from warden.lib.models import Job, Session +# Deliberately generous: this budget is only ever spent by a test that is +# already failing, so it costs nothing on the happy path. Tight per-test budgets +# turned a slow CI runner into a flake instead of catching anything. +JOB_WAIT_TIMEOUT_S = 30 + async def wait_until_scalar_equals( session: AsyncSession, @@ -123,6 +130,31 @@ async def scheduler_task_timeout(delay: float, scheduler_task: Task): pass +async def wait_until_jobs_settled( + session: AsyncSession, + scheduler_task: Task, + *, + count: int, + statuses: Sequence[str] = ("DONE",), + timeout_s: float = JOB_WAIT_TIMEOUT_S, + interval: float = 0.1, +) -> None: + """Wait until ``count`` jobs reached one of ``statuses``, then stop the scheduler. + + Use this rather than hand-rolling a wait predicate. The scheduler commits a + job's terminal status and its complete logs in a single transaction (see + ``JobExecutionTracker.update_job``), so waiting on the status alone is enough + to make the assertions that follow - including any on ``logs`` - safe. + + Fails the test on timeout, and always cancels ``scheduler_task`` so it cannot + outlive the test body and interfere with fixture teardown. + """ + + stmt = select(func.count(Job.id)).where(Job.status.in_(tuple(statuses))) + async with scheduler_task_timeout(timeout_s, scheduler_task): + await wait_until_scalar_equals(session, stmt, count, interval=interval) + + def build_conf(strategy: SchedulerStrategy, qpu_uri: str) -> Config: return Config( scheduler=SchedulerConfig( diff --git a/tests/test_config.py b/tests/test_config.py index 912c34b..ba2cedd 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,6 +8,8 @@ from warden.lib.config.config import ( APIConfig, Config, + QPUAuthConfig, + QPUConfig, SchedulerConfig, SchedulerStrategy, ) @@ -106,3 +108,77 @@ def test_admin_users_must_not_be_empty(): """ with pytest.raises(ValidationError): APIConfig(admin_users=[]) + + +def test_qpu_auth_absent_by_default(): + assert Config().qpu.auth is None + + +def test_qpu_auth_token_url_is_built_from_base_and_realm(): + auth = QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_token_url_tolerates_trailing_slash(): + auth = QPUAuthConfig( + url="http://keycloak:8080/", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_rejects_partial_configuration(): + # A half-configured auth section must fail loudly rather than silently + # falling back to unauthenticated requests. + with pytest.raises(ValidationError): + # model_validate, not the constructor: omitting a required field is the + # point of the test, and a static type checker rejects the direct call. + QPUAuthConfig.model_validate({"url": "http://keycloak:8080", "id": "warden"}) + + +def test_qpu_auth_secret_read_from_env(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("WARDEN_QPU_AUTH_URL", "http://keycloak:8080") + monkeypatch.setenv("WARDEN_QPU_AUTH_ID", "warden") + monkeypatch.setenv("WARDEN_QPU_AUTH_SECRET", "from-env") + + config = Config() + + assert config.qpu.auth is not None + assert config.qpu.auth.id == "warden" + assert config.qpu.auth.secret == "from-env" + + +def test_auth_flow_is_none_without_auth_config(): + assert Config().qpu.auth_flow is None + + +def test_auth_flow_is_memoized(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.auth_flow is qpu.auth_flow + + +def test_client_is_given_the_auth_flow(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.client.auth is qpu.auth_flow + + +def test_client_has_no_auth_without_auth_config(): + assert QPUConfig(uri="http://qpu:4300").client.auth is None diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index a9fc30d..191a691 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -70,8 +70,37 @@ class SchedulerConfig(BaseSettings): job_polling_timeout_s: float = -1 +class QPUAuthConfig(BaseSettings): + """Keycloak client_credentials configuration for outbound QPU API calls. + + Presence of this section is what enables authentication. There is + deliberately no separate ``enabled`` flag: a second switch can drift out of + sync with the credentials it guards. ``url``, ``id`` and ``secret`` have no + defaults, so a partially configured section is a startup validation error + rather than a silent fallback to unauthenticated requests. + """ + + # Keycloak base URL, e.g. http://keycloak:8080 + url: str + realm: str = "pasqos" + # OIDC client_id + id: str + # OIDC client_secret. Provide via WARDEN_QPU_AUTH_SECRET, never in YAML. + secret: str + # Refresh this many seconds before the token actually expires. + leeway_s: float = 30 + + @property + def token_url(self) -> str: + """Keycloak's OIDC token endpoint for this realm.""" + return ( + f"{self.url.rstrip('/')}/realms/{self.realm}/protocol/openid-connect/token" + ) + + class QPUConfig(BaseSettings): uri: str = "http://localhost:8000" + auth: QPUAuthConfig | None = None retry_max: int = 10 retry_sleep_s: float = 1 @@ -86,6 +115,7 @@ class QPUConfig(BaseSettings): tls_verify: bool | str = "system" _client: httpx.Client | None = PrivateAttr(default=None) + _auth_flow: httpx.Auth | None = PrivateAttr(default=None) @property def verify(self) -> bool | str | ssl.SSLContext: @@ -96,10 +126,28 @@ def verify(self) -> bool | str | ssl.SSLContext: return ssl.create_default_context() return self.tls_verify + @property + def auth_flow(self) -> httpx.Auth | None: + """Memoized Keycloak auth flow, or None when auth is not configured. + + Memoized so the sync (scheduler) and async (API) clients in one process + share a single cached token. Imported lazily because + ``qpu_client.auth`` imports this module. + """ + if self.auth is None: + return None + if self._auth_flow is None: + from warden.lib.qpu_client.auth import KeycloakClientCredentialsAuth + + self._auth_flow = KeycloakClientCredentialsAuth( + self.auth, verify=self.verify + ) + return self._auth_flow + @property def client(self) -> httpx.Client: if self._client is None: - self._client = httpx.Client(verify=self.verify) + self._client = httpx.Client(verify=self.verify, auth=self.auth_flow) self._client.base_url = self.uri + API_PREFIX return self._client diff --git a/warden/lib/config/config.sample.yaml b/warden/lib/config/config.sample.yaml index 0b35b20..10a8512 100644 --- a/warden/lib/config/config.sample.yaml +++ b/warden/lib/config/config.sample.yaml @@ -73,6 +73,21 @@ scheduler: qpu: # Local Pasqal QPU API configuration uri: http://localhost:8000 + # Keycloak authentication for outbound requests to the QPU API. + # Omit this whole section to send unauthenticated requests (e.g. against the + # mock QPU). If present, 'url', 'id' and 'secret' are all mandatory: a + # half-configured section is a startup error rather than a silent fallback + # to unauthenticated requests. + # auth: + # # Keycloak base URL (not the token endpoint; the OIDC path is appended). + # url: http://keycloak:8080 + # realm: pasqos + # # OIDC client_id of the service account. + # id: warden + # # Never put the secret in this file. Set it in the environment: + # # WARDEN_QPU_AUTH_SECRET="..." + # # Refresh the token this many seconds before it expires. + # leeway_s: 30 # TLS verification policy for requests to the QPU backend. # Only relevant when 'uri' uses https. Accepts: diff --git a/warden/lib/qpu_client/auth.py b/warden/lib/qpu_client/auth.py new file mode 100644 index 0000000..046ec26 --- /dev/null +++ b/warden/lib/qpu_client/auth.py @@ -0,0 +1,138 @@ +"""Keycloak client_credentials authentication for outbound QPU API calls.""" + +import logging +import ssl +from time import monotonic +from typing import AsyncGenerator, Generator + +import httpx + +from warden.lib.config.config import QPUAuthConfig +from warden.lib.qpu_client.retry import QPUClientRequestError + +logger = logging.getLogger(__name__) + +# Token-endpoint statuses that will never succeed on retry: the credentials or +# the grant itself are wrong. Anything else (transport errors, 5xx) is left to +# propagate so the existing retry decorator can treat it as transient. +FATAL_TOKEN_STATUSES = (400, 401, 403) + + +class TokenRequestError(QPUClientRequestError): + """Keycloak refused to issue a token and retrying cannot help.""" + + +class KeycloakClientCredentialsAuth(httpx.Auth): + """Attach a Keycloak service-account bearer token to each request. + + Implemented as an ``httpx.Auth`` so it runs inside the transport, below + Warden's ``retry`` decorator. That matters because 401 is not in + ``RETRY_HTTP_EXIT_CODES``: a token expiring mid-job would otherwise surface + as an immediate, non-retryable ``NotRetriedHTTPStatus``. Here it is just a + refresh. + + Args: + conf: Keycloak credentials and endpoint. + verify: httpx TLS verification setting for the token request. + """ + + def __init__( + self, + conf: QPUAuthConfig, + verify: bool | str | ssl.SSLContext = True, + ) -> None: + self.conf = conf + self.verify = verify + self._token: str | None = None + # monotonic() deadline after which the cached token is considered stale. + self._expires_at: float = 0.0 + + def sync_auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + request.headers["Authorization"] = f"Bearer {self._sync_token()}" + response = yield request + if response.status_code == httpx.codes.UNAUTHORIZED: + logger.info("QPU API returned 401, refreshing token and retrying once") + request.headers["Authorization"] = f"Bearer {self._sync_token(force=True)}" + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + request.headers["Authorization"] = f"Bearer {await self._async_token()}" + response = yield request + if response.status_code == httpx.codes.UNAUTHORIZED: + logger.info("QPU API returned 401, refreshing token and retrying once") + request.headers["Authorization"] = ( + f"Bearer {await self._async_token(force=True)}" + ) + yield request + + def _is_fresh(self) -> bool: + return self._token is not None and monotonic() < self._expires_at + + def _token_request(self) -> tuple[str, dict[str, str]]: + """Return the (url, form data) for a client_credentials token request.""" + return self.conf.token_url, { + "grant_type": "client_credentials", + "client_id": self.conf.id, + "client_secret": self.conf.secret, + } + + def _store(self, response: httpx.Response) -> str: + """Validate a token response, cache the token and return it.""" + if response.status_code in FATAL_TOKEN_STATUSES: + # Never log the response body of a token request: it may echo + # credentials. The error field alone is the useful part. + try: + error = response.json().get("error", "unknown_error") + except (ValueError, AttributeError): + error = "unknown_error" + raise TokenRequestError( + f"Keycloak refused to issue a token for client " + f"'{self.conf.id}' at {self.conf.token_url}: " + f"{response.status_code} {error}" + ) + # Transport errors and 5xx stay as httpx exceptions so the existing + # retry decorator sees them as transient. + response.raise_for_status() + + payload = response.json() + token = payload["access_token"] + expires_in = float(payload.get("expires_in", 0)) + self._token = token + ttl = max(expires_in - self.conf.leeway_s, expires_in / 2) + if expires_in <= self.conf.leeway_s: + logger.warning( + f"Token lifespan {expires_in}s is not greater than the " + f"configured leeway {self.conf.leeway_s}s; caching for {ttl}s " + "(half the lifespan) instead of refusing to cache at all" + ) + self._expires_at = monotonic() + ttl + logger.debug( + f"Obtained QPU API token for client '{self.conf.id}', " + f"expires in {expires_in}s" + ) + return token + + def _sync_token(self, force: bool = False) -> str: + if not force and self._is_fresh(): + assert self._token is not None + return self._token + url, data = self._token_request() + with httpx.Client(verify=self.verify) as client: + return self._store(client.post(url, data=data)) + + # Note: unlocked check-then-fetch. Two concurrent async requests in one + # process can both miss and both fetch a token; one wins and the loser + # wasted a request. The sync path (scheduler) blocks its single event loop + # per fetch, so this race is only reachable via awaited callers here. Add + # a lock only if token-endpoint traffic ever becomes a problem. + async def _async_token(self, force: bool = False) -> str: + if not force and self._is_fresh(): + assert self._token is not None + return self._token + url, data = self._token_request() + async with httpx.AsyncClient(verify=self.verify) as client: + return self._store(await client.post(url, data=data)) diff --git a/warden/lib/qpu_client/client.py b/warden/lib/qpu_client/client.py index 9483206..847c781 100644 --- a/warden/lib/qpu_client/client.py +++ b/warden/lib/qpu_client/client.py @@ -198,7 +198,9 @@ class AsyncQPUClient: def __init__(self, qpu_conf: QPUConfig): self.conf = qpu_conf self.client = AsyncClient( - base_url=qpu_conf.uri + "/api/v1", verify=qpu_conf.verify + base_url=qpu_conf.uri + "/api/v1", + verify=qpu_conf.verify, + auth=qpu_conf.auth_flow, ) async def get_specs(self) -> str: diff --git a/warden/lib/qpu_client/retry.py b/warden/lib/qpu_client/retry.py index 30b61a6..2eaf1e8 100644 --- a/warden/lib/qpu_client/retry.py +++ b/warden/lib/qpu_client/retry.py @@ -52,6 +52,8 @@ def retry(max: int, sleep_s: float, no_retry: bool = False) -> Callable: UnhandledError: If decorator encounters an unnexpected exception. NotRetriedHTTPStatus: If the HTTP request returns with a non-retryable error code. MaxRetryError: If the maximum number of retries without success has been reached. + QPUClientRequestError: Any subclass already classified as non-retryable + by the wrapped function (e.g. TokenRequestError) propagates unchanged. """ def decorator(func: Callable): @@ -62,6 +64,10 @@ def _handle_exception(e: Exception): elif isinstance(e, HTTPStatusError): if e.response.status_code not in RETRY_HTTP_EXIT_CODES: raise NotRetriedHTTPStatus(e) from e + elif isinstance(e, QPUClientRequestError): + # Already classified as non-retryable by the raiser (e.g. bad + # Keycloak credentials). Do not rewrap it as UnhandledError. + raise else: raise UnhandledError(e) from e diff --git a/warden/scheduler/worker.py b/warden/scheduler/worker.py index a29a2a1..eb9f9b5 100644 --- a/warden/scheduler/worker.py +++ b/warden/scheduler/worker.py @@ -19,6 +19,8 @@ logger = logging.getLogger(__name__) +TERMINAL_STATUSES: tuple[JobStatus, ...] = ("ERROR", "DONE", "CANCELED") + class JobExecutionTracker: """Handles current job status and sends updates to db""" @@ -49,8 +51,16 @@ def created_datetime(self) -> UTCDatetime: return self.job.created_datetime async def update_job(self, qpu_job_info: QPUJobInfo): + was_terminal = self._status in TERMINAL_STATUSES self._qpu_job_info = qpu_job_info self._status = qpu_job_info.status or "ERROR" + if self._status in TERMINAL_STATUSES and not was_terminal: + # Logged here rather than by the caller so that the closing line is + # part of the same JobUpdate - hence the same transaction - as the + # terminal status. Flushed separately, the DB would briefly hold a + # finished job whose logs are truncated, and anything that stops + # polling once the status is terminal reads incomplete logs. + logger.info("Job execution ended with status '%s'", self._status) await self.push_update() async def to_error(self): @@ -166,7 +176,6 @@ async def execute_job( return await self.await_job_execution(job_tracker) - logger.info("Job execution ended with status '%s'", job_tracker.status) # Flush potential last updates before return await job_tracker.push_update() @@ -241,7 +250,7 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: polling_start = job_tracker.created_datetime await self._get_job_poll(job_tracker) - while job_tracker.status not in ("ERROR", "DONE", "CANCELED"): + while job_tracker.status not in TERMINAL_STATUSES: if self.is_timed_out(self.conf_sched.job_polling_timeout_s, polling_start): logger.warning( f"Job timed out (max {self.conf_sched.job_polling_timeout_s} s). " @@ -250,12 +259,15 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: ) try: qpu_job_info = self.qpu_client.cancel_job(job_tracker.job.uid) + # Logged before the update so it is buffered into the same + # JobUpdate, and stays ahead of the closing line that a + # terminal status appends + logger.info("Job cancellation done") await job_tracker.update_job(qpu_job_info) except (JobCancelationError, QPUClientRequestError) as e: logger.error(f"Failed cancelling job: {e}") await job_tracker.to_error() continue - logger.info("Job cancellation done") continue await asyncio.sleep(self.conf_sched.job_polling_interval_s) await self._get_job_poll(job_tracker)