mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
Shahar Epstein
199b0a8552
CI no longer runs Python 3.10, so the tomli fallbacks for tomllib, the typing_extensions imports of Self, Unpack and assert_never, the re-implementation of asyncio.Runner in the Task SDK and the version gated test skips no longer have a version to serve.
2378 lines
100 KiB
Python
2378 lines
100 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import pickle
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING
|
|
from unittest import mock
|
|
|
|
import certifi
|
|
import httpx
|
|
import pytest
|
|
import time_machine
|
|
import uuid6
|
|
from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses
|
|
from uuid6 import uuid7
|
|
|
|
from airflow.sdk import timezone
|
|
from airflow.sdk.api.client import Client, RemoteValidationError, ServerResponseError
|
|
from airflow.sdk.api.datamodels._generated import (
|
|
AssetEventsResponse,
|
|
AssetResponse,
|
|
AssetStateStoreResponse,
|
|
ConnectionResponse,
|
|
DagResponse,
|
|
DagRunState,
|
|
DagRunStateResponse,
|
|
HITLDetailRequest,
|
|
HITLDetailResponse,
|
|
HITLUser,
|
|
TaskStateStoreResponse,
|
|
TerminalTIState,
|
|
VariableResponse,
|
|
XComResponse,
|
|
)
|
|
from airflow.sdk.exceptions import ErrorType, TaskAlreadyRunningError
|
|
from airflow.sdk.execution_time.comms import (
|
|
AssetsByAliasResult,
|
|
DeferTask,
|
|
ErrorResponse,
|
|
OKResponse,
|
|
PreviousDagRunResult,
|
|
PreviousTIResult,
|
|
RescheduleTask,
|
|
TaskRescheduleStartDate,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from time_machine import TimeMachineFixture
|
|
|
|
|
|
class TestClient:
|
|
@pytest.mark.parametrize(
|
|
("path", "json_response"),
|
|
[
|
|
(
|
|
"/task-instances/1/run",
|
|
{
|
|
"dag_run": {
|
|
"dag_id": "test_dag",
|
|
"run_id": "test_run",
|
|
"logical_date": "2021-01-01T00:00:00Z",
|
|
"start_date": "2021-01-01T00:00:00Z",
|
|
"run_type": "manual",
|
|
"run_after": "2021-01-01T00:00:00Z",
|
|
"consumed_asset_events": [],
|
|
},
|
|
"max_tries": 0,
|
|
"should_retry": False,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_dry_run(self, path, json_response):
|
|
client = make_client_w_dry_run()
|
|
assert client.base_url == "dry-run://server"
|
|
|
|
resp = client.get(path)
|
|
|
|
assert resp.status_code == 200
|
|
assert resp.json() == json_response
|
|
|
|
@mock.patch("airflow.sdk.api.client.API_SSL_CERT_PATH", "/capath/does/not/exist/")
|
|
def test_add_capath(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
with pytest.raises(FileNotFoundError) as err:
|
|
make_client(httpx.MockTransport(handle_request))
|
|
|
|
assert isinstance(err.value, FileNotFoundError)
|
|
|
|
@mock.patch("airflow.sdk.api.client.API_SSL_CA_FILE_PATH", "/capath/does/not/exist/")
|
|
def test_ssl_with_cafile(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
with pytest.raises(FileNotFoundError) as err:
|
|
make_client(httpx.MockTransport(handle_request))
|
|
|
|
assert isinstance(err.value, FileNotFoundError)
|
|
|
|
@mock.patch("ssl.create_default_context")
|
|
@mock.patch("airflow.sdk.api.client.API_CLIENT_USE_PUBLIC_CERTS", True)
|
|
def test_use_public_certs(self, mock_default_context):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
make_client(httpx.MockTransport(handle_request))
|
|
mock_default_context.return_value.load_verify_locations.assert_called_with(certifi.where())
|
|
|
|
@mock.patch("airflow.sdk.api.client.API_TIMEOUT", 60.0)
|
|
def test_timeout_configuration(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
client = make_client(httpx.MockTransport(handle_request))
|
|
assert client.timeout == httpx.Timeout(60.0)
|
|
|
|
def test_timeout_can_be_overridden(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
client = Client(
|
|
base_url="test://server", token="", transport=httpx.MockTransport(handle_request), timeout=120.0
|
|
)
|
|
assert client.timeout == httpx.Timeout(120.0)
|
|
|
|
def test_error_parsing(self):
|
|
responses = [
|
|
httpx.Response(422, json={"detail": [{"loc": ["#0"], "msg": "err", "type": "required"}]})
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(ServerResponseError) as err:
|
|
client.get("http://error")
|
|
|
|
assert isinstance(err.value, ServerResponseError)
|
|
assert isinstance(err.value.detail, list)
|
|
assert err.value.detail == [
|
|
RemoteValidationError(loc=["#0"], msg="err", type="required"),
|
|
]
|
|
|
|
def test_error_parsing_plain_text(self):
|
|
responses = [httpx.Response(422, content=b"Internal Server Error")]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(httpx.HTTPStatusError) as err:
|
|
client.get("http://error")
|
|
assert not isinstance(err.value, ServerResponseError)
|
|
|
|
def test_error_parsing_other_json(self):
|
|
responses = [httpx.Response(404, json={"detail": "Not found"})]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(ServerResponseError) as err:
|
|
client.get("http://error")
|
|
assert err.value.args == ("Not found",)
|
|
assert err.value.detail is None
|
|
|
|
def test_server_response_error_pickling(self):
|
|
responses = [httpx.Response(404, json={"detail": {"message": "Invalid input"}})]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(ServerResponseError) as exc_info:
|
|
client.get("http://error")
|
|
|
|
err = exc_info.value
|
|
assert err.args == ("Server returned error",)
|
|
assert err.detail == {"message": "Invalid input"}
|
|
|
|
# Check that the error is picklable
|
|
pickled = pickle.dumps(err)
|
|
unpickled = pickle.loads(pickled)
|
|
|
|
assert isinstance(unpickled, ServerResponseError)
|
|
|
|
# Test that unpickled error has the same attributes as the original
|
|
assert unpickled.response.json() == {"detail": {"message": "Invalid input"}}
|
|
assert unpickled.detail == {"message": "Invalid input"}
|
|
assert unpickled.response.status_code == 404
|
|
assert unpickled.request.url == "http://error"
|
|
|
|
def test_server_error_detail_added_as_note(self):
|
|
"""Notes survive uncaught propagation, handled sites still log detail directly."""
|
|
responses = [httpx.Response(404, json={"detail": {"message": "Invalid input"}})]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(ServerResponseError) as exc_info:
|
|
client.get("http://error")
|
|
|
|
err = exc_info.value
|
|
assert err.args == ("Server returned error",)
|
|
assert any(
|
|
note.startswith("Server error detail:") and "Invalid input" in note
|
|
for note in getattr(err, "__notes__", [])
|
|
), err.__notes__
|
|
|
|
def test_retry_handling_unrecoverable_error(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
responses: list[httpx.Response] = [
|
|
*[httpx.Response(500, text="Internal Server Error")] * 6,
|
|
httpx.Response(200, json={"detail": "Recovered from error - but will fail before"}),
|
|
httpx.Response(400, json={"detail": "Should not get here"}),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(httpx.HTTPStatusError) as err:
|
|
client.get("http://error")
|
|
assert not isinstance(err.value, ServerResponseError)
|
|
assert len(responses) == 3
|
|
|
|
def test_retry_handling_recovered(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
responses: list[httpx.Response] = [
|
|
*[httpx.Response(500, text="Internal Server Error")] * 2,
|
|
httpx.Response(200, json={"detail": "Recovered from error"}),
|
|
httpx.Response(400, json={"detail": "Should not get here"}),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
response = client.get("http://error")
|
|
assert response.status_code == 200
|
|
assert len(responses) == 1
|
|
|
|
def test_retry_handling_non_retry_error(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
responses: list[httpx.Response] = [
|
|
httpx.Response(422, json={"detail": "Somehow this is a bad request"}),
|
|
httpx.Response(400, json={"detail": "Should not get here"}),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(ServerResponseError) as err:
|
|
client.get("http://error")
|
|
assert len(responses) == 1
|
|
assert err.value.args == ("Somehow this is a bad request",)
|
|
|
|
def test_retry_handling_ok(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
responses: list[httpx.Response] = [
|
|
httpx.Response(200, json={"detail": "Recovered from error"}),
|
|
httpx.Response(400, json={"detail": "Should not get here"}),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
response = client.get("http://error")
|
|
assert response.status_code == 200
|
|
assert len(responses) == 1
|
|
|
|
def test_token_renewal(self):
|
|
responses: list[httpx.Response] = [
|
|
httpx.Response(200, json={"ok": "1"}),
|
|
httpx.Response(404, json={"var": "not_found"}, headers={"Refreshed-API-Token": "abc"}),
|
|
httpx.Response(200, json={"ok": "3"}),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert client.auth is not None
|
|
assert not client.auth.token
|
|
with pytest.raises(ServerResponseError):
|
|
response = client.get("/")
|
|
|
|
# Even thought it was Not Found, we should still respect the header
|
|
assert client.auth is not None
|
|
assert client.auth.token == "abc"
|
|
|
|
# Test that the next request is made with that new auth token
|
|
response = client.get("/")
|
|
assert response.status_code == 200
|
|
assert response.request.headers["Authorization"] == "Bearer abc"
|
|
|
|
@pytest.mark.parametrize(
|
|
("status_code", "description"),
|
|
[
|
|
(399, "status code < 400"),
|
|
(301, "3xx redirect status code"),
|
|
(600, "status code >= 600"),
|
|
],
|
|
)
|
|
def test_server_response_error_invalid_status_codes(self, status_code, description):
|
|
"""Test that ServerResponseError.from_response returns None for invalid status codes."""
|
|
response = httpx.Response(status_code, json={"detail": f"Test {description}"})
|
|
assert ServerResponseError.from_response(response) is None
|
|
|
|
@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_CERT", "/etc/airflow/certs/client.crt")
|
|
@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_KEY", "/etc/airflow/certs/client.key")
|
|
def test_sets_cert_tuple(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
captured: dict[str, object] = {}
|
|
real_init = httpx.Client.__init__
|
|
|
|
def spy_init(self, *args, **kwargs):
|
|
captured["cert"] = kwargs.get("cert")
|
|
return real_init(self, *args, **kwargs)
|
|
|
|
with mock.patch.object(httpx.Client, "__init__", spy_init):
|
|
make_client(httpx.MockTransport(handle_request))
|
|
|
|
assert captured["cert"] == ("/etc/airflow/certs/client.crt", "/etc/airflow/certs/client.key")
|
|
|
|
@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_CERT", "/etc/airflow/certs/client.crt")
|
|
@mock.patch("airflow.sdk.api.client.API_CLIENT_SSL_KEY", None)
|
|
def test_requires_both_cert_and_key(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200)
|
|
|
|
with pytest.raises(ValueError, match="Both client_ssl_cert and client_ssl_key must be set"):
|
|
make_client(httpx.MockTransport(handle_request))
|
|
|
|
|
|
class TestTaskInstanceOperations:
|
|
"""
|
|
Test that the TestTaskInstanceOperations class works as expected. While the operations are simple, it
|
|
still catches the basic functionality of the client for task instances including endpoint and
|
|
response parsing.
|
|
"""
|
|
|
|
def test_task_instance_start(self, make_ti_context):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
# Simulate a successful response from the server that starts a task
|
|
ti_id = uuid6.uuid7()
|
|
start_date = "2024-10-31T12:00:00Z"
|
|
ti_context = make_ti_context(
|
|
start_date=start_date,
|
|
logical_date="2024-10-31T12:00:00Z",
|
|
run_type="manual",
|
|
)
|
|
|
|
# ...including a validation that retry really works
|
|
call_count = 0
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count < 3:
|
|
return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
|
|
if request.url.path == f"/task-instances/{ti_id}/run":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["pid"] == 100
|
|
assert actual_body["start_date"] == start_date
|
|
assert actual_body["state"] == "running"
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json=ti_context.model_dump(mode="json"),
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
resp = client.task_instances.start(ti_id, 100, start_date)
|
|
assert resp == ti_context
|
|
assert call_count == 3
|
|
|
|
@pytest.mark.parametrize(
|
|
("reason", "previous_state"),
|
|
[("invalid_state", "running"), ("running_elsewhere", "restarting")],
|
|
)
|
|
def test_task_instance_start_already_running(self, reason, previous_state):
|
|
"""Test that start() raises TaskAlreadyRunningError when TI is already running."""
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/run":
|
|
return httpx.Response(
|
|
409,
|
|
json={
|
|
"detail": {
|
|
"reason": reason,
|
|
"message": "TI was not in a state where it could be marked as running",
|
|
"previous_state": previous_state,
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with pytest.raises(TaskAlreadyRunningError, match="already running"):
|
|
client.task_instances.start(ti_id, 100, datetime(2024, 10, 31, tzinfo=timezone.utc))
|
|
|
|
@pytest.mark.parametrize("previous_state", ["failed", "success", "skipped"])
|
|
def test_task_instance_start_other_invalid_states(self, previous_state):
|
|
"""Test that start() raises ServerResponseError for non-running invalid states."""
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/run":
|
|
return httpx.Response(
|
|
409,
|
|
json={
|
|
"detail": {
|
|
"reason": "invalid_state",
|
|
"message": "TI was not in a state where it could be marked as running",
|
|
"previous_state": previous_state,
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with pytest.raises(ServerResponseError):
|
|
client.task_instances.start(ti_id, 100, datetime(2024, 10, 31, tzinfo=timezone.utc))
|
|
|
|
@pytest.mark.parametrize(
|
|
"state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS]
|
|
)
|
|
@pytest.mark.parametrize("retry_reason", [None, "auth error, do not retry"])
|
|
def test_task_instance_finish(self, state, retry_reason):
|
|
# Simulate a successful response from the server that finishes (moved to terminal state) a task
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/state":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["end_date"] == "2024-10-31T12:00:00Z"
|
|
assert actual_body["state"] == state
|
|
assert actual_body["rendered_map_index"] == "test"
|
|
assert actual_body["retry_reason"] == retry_reason
|
|
return httpx.Response(
|
|
status_code=204,
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_instances.finish(
|
|
ti_id,
|
|
state=state,
|
|
when="2024-10-31T12:00:00Z",
|
|
rendered_map_index="test",
|
|
retry_reason=retry_reason,
|
|
)
|
|
|
|
def test_task_instance_heartbeat(self):
|
|
# Simulate a successful response from the server that sends a heartbeat for a ti
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/heartbeat":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["pid"] == 100
|
|
return httpx.Response(
|
|
status_code=204,
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_instances.heartbeat(ti_id, 100)
|
|
|
|
def test_task_instance_heartbeat_is_not_retried(self):
|
|
"""A failing heartbeat surfaces at once so the supervisor's own retry budget governs it."""
|
|
ti_id = uuid6.uuid7()
|
|
responses: list[httpx.Response] = [
|
|
httpx.Response(500, text="Internal Server Error"),
|
|
httpx.Response(204),
|
|
]
|
|
client = make_client_w_responses(responses)
|
|
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
client.task_instances.heartbeat(ti_id, 100)
|
|
|
|
assert len(responses) == 1
|
|
|
|
def test_task_instance_update_dagrun_note(self):
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/dag-run-note":
|
|
assert json.loads(request.read()) == {"note": "Updated from task runtime"}
|
|
return httpx.Response(status_code=204)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
response = client.task_instances.update_dagrun_note(ti_id, "Updated from task runtime")
|
|
|
|
assert response == OKResponse(ok=True)
|
|
|
|
@pytest.mark.parametrize("queues_enabled", [False, True])
|
|
def test_task_instance_defer(self, queues_enabled: bool):
|
|
# Simulate a successful response from the server that defers a task
|
|
ti_id = uuid6.uuid7()
|
|
task_queue = "test"
|
|
|
|
msg = DeferTask(
|
|
classpath="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
|
|
next_method="execute_complete",
|
|
trigger_kwargs={
|
|
"__type": "dict",
|
|
"__var": {
|
|
"moment": {"__type": "datetime", "__var": 1730982899.0},
|
|
"end_from_trigger": False,
|
|
},
|
|
},
|
|
next_kwargs={"__type": "dict", "__var": {}},
|
|
queue=task_queue if queues_enabled else None,
|
|
)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/state":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["state"] == "deferred"
|
|
assert actual_body["trigger_kwargs"] == msg.trigger_kwargs
|
|
assert (
|
|
actual_body["classpath"] == "airflow.providers.standard.triggers.temporal.DateTimeTrigger"
|
|
)
|
|
assert actual_body["next_method"] == "execute_complete"
|
|
if queues_enabled:
|
|
assert actual_body["queue"] == task_queue
|
|
else:
|
|
assert actual_body.get("queue") is None
|
|
return httpx.Response(
|
|
status_code=204,
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_instances.defer(ti_id, msg)
|
|
|
|
def test_task_instance_reschedule(self):
|
|
# Simulate a successful response from the server that reschedules a task
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/state":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["state"] == "up_for_reschedule"
|
|
assert actual_body["reschedule_date"] == "2024-10-31T12:00:00Z"
|
|
assert actual_body["end_date"] == "2024-10-31T12:00:00Z"
|
|
return httpx.Response(
|
|
status_code=204,
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
msg = RescheduleTask(
|
|
reschedule_date=timezone.parse("2024-10-31T12:00:00Z"),
|
|
end_date=timezone.parse("2024-10-31T12:00:00Z"),
|
|
)
|
|
client.task_instances.reschedule(ti_id, msg)
|
|
|
|
def test_task_instance_up_for_retry(self):
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{ti_id}/state":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["state"] == "up_for_retry"
|
|
assert actual_body["end_date"] == "2024-10-31T12:00:00Z"
|
|
assert actual_body["rendered_map_index"] == "test"
|
|
return httpx.Response(
|
|
status_code=204,
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_instances.retry(
|
|
ti_id, end_date=timezone.parse("2024-10-31T12:00:00Z"), rendered_map_index="test"
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"rendered_fields",
|
|
[
|
|
pytest.param({"field1": "rendered_value1", "field2": "rendered_value2"}, id="simple-rendering"),
|
|
pytest.param(
|
|
{
|
|
"field1": "ClassWithCustomAttributes({'nested1': ClassWithCustomAttributes("
|
|
"{'att1': 'test', 'att2': 'test2'), "
|
|
"'nested2': ClassWithCustomAttributes("
|
|
"{'att3': 'test3', 'att4': 'test4')"
|
|
},
|
|
id="complex-rendering",
|
|
),
|
|
],
|
|
)
|
|
def test_taskinstance_set_rtif_success(self, rendered_fields):
|
|
TI_ID = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{TI_ID}/rtif":
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"message": "Rendered task instance fields successfully set"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.set_rtif(id=TI_ID, body=rendered_fields)
|
|
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_taskinstance_set_rendered_map_index_success(self):
|
|
TI_ID = uuid6.uuid7()
|
|
rendered_map_index = "Label: task_1"
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-instances/{TI_ID}/rendered-map-index":
|
|
actual_body = json.loads(request.read())
|
|
assert request.method == "PATCH"
|
|
# Body should be the string directly, not wrapped in JSON
|
|
assert actual_body == rendered_map_index
|
|
return httpx.Response(status_code=204)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.set_rendered_map_index(id=TI_ID, rendered_map_index=rendered_map_index)
|
|
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_get_count_basic(self):
|
|
"""Test basic get_count functionality with just dag_id."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/task-instances/count"
|
|
assert request.url.params.get("dag_id") == "test_dag"
|
|
return httpx.Response(200, json=5)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_count(dag_id="test_dag")
|
|
assert result.count == 5
|
|
|
|
def test_get_count_with_all_params(self):
|
|
"""Test get_count with all optional parameters."""
|
|
|
|
logical_dates_str = ["2024-01-01T00:00:00+00:00", "2024-01-02T00:00:00+00:00"]
|
|
logical_dates = [timezone.parse(d) for d in logical_dates_str]
|
|
task_ids = ["task1", "task2"]
|
|
states = ["success", "failed"]
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/task-instances/count"
|
|
assert request.method == "GET"
|
|
params = request.url.params
|
|
assert params["dag_id"] == "test_dag"
|
|
assert params.get_list("task_ids") == task_ids
|
|
assert params["task_group_id"] == "group1"
|
|
assert params.get_list("logical_dates") == logical_dates_str
|
|
assert params.get_list("run_ids") == []
|
|
assert params.get_list("states") == states
|
|
assert params["map_index"] == "0"
|
|
return httpx.Response(200, json=10)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_count(
|
|
dag_id="test_dag",
|
|
map_index=0,
|
|
task_ids=task_ids,
|
|
task_group_id="group1",
|
|
logical_dates=logical_dates,
|
|
states=states,
|
|
)
|
|
assert result.count == 10
|
|
|
|
def test_get_task_states_basic(self):
|
|
"""Test basic get_task_states functionality with just dag_id."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/task-instances/states"
|
|
assert request.url.params.get("dag_id") == "test_dag"
|
|
assert request.url.params.get("task_group_id") == "group1"
|
|
return httpx.Response(
|
|
200, json={"task_states": {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}}
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_task_states(dag_id="test_dag", task_group_id="group1")
|
|
assert result.task_states == {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}
|
|
|
|
def test_get_task_states_with_all_params(self):
|
|
"""Test get_task_states with all optional parameters."""
|
|
|
|
logical_dates_str = ["2024-01-01T00:00:00+00:00", "2024-01-02T00:00:00+00:00"]
|
|
logical_dates = [timezone.parse(d) for d in logical_dates_str]
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/task-instances/states"
|
|
assert request.method == "GET"
|
|
params = request.url.params
|
|
assert params["dag_id"] == "test_dag"
|
|
assert params["task_group_id"] == "group1"
|
|
assert params.get_list("logical_dates") == logical_dates_str
|
|
assert params.get_list("task_ids") == []
|
|
assert params.get_list("run_ids") == []
|
|
assert params.get("map_index") == "0"
|
|
return httpx.Response(
|
|
200, json={"task_states": {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}}
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_task_states(
|
|
dag_id="test_dag",
|
|
map_index=0,
|
|
task_group_id="group1",
|
|
logical_dates=logical_dates,
|
|
)
|
|
assert result.task_states == {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}
|
|
|
|
def test_get_previous_basic(self):
|
|
"""Test basic get_previous functionality."""
|
|
logical_date = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/task-instances/previous/test_dag/test_task":
|
|
assert request.url.params == httpx.QueryParams(
|
|
logical_date=logical_date.isoformat(),
|
|
map_index="-1",
|
|
)
|
|
# Return complete TI data
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"task_id": "test_task",
|
|
"dag_id": "test_dag",
|
|
"run_id": "prev_run",
|
|
"logical_date": "2024-01-14T12:00:00+00:00",
|
|
"start_date": "2024-01-14T12:05:00+00:00",
|
|
"end_date": "2024-01-14T12:10:00+00:00",
|
|
"state": "success",
|
|
"try_number": 1,
|
|
"map_index": -1,
|
|
"duration": 300.0,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_previous(
|
|
dag_id="test_dag", task_id="test_task", logical_date=logical_date
|
|
)
|
|
|
|
assert isinstance(result, PreviousTIResult)
|
|
assert result.task_instance.task_id == "test_task"
|
|
assert result.task_instance.dag_id == "test_dag"
|
|
assert result.task_instance.run_id == "prev_run"
|
|
assert result.task_instance.state == "success"
|
|
|
|
def test_get_previous_with_state_filter(self):
|
|
"""Test get_previous functionality with state filtering."""
|
|
logical_date = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/task-instances/previous/test_dag/test_task":
|
|
assert request.url.params == httpx.QueryParams(
|
|
logical_date=logical_date.isoformat(),
|
|
map_index="-1",
|
|
state="success",
|
|
)
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"task_id": "test_task",
|
|
"dag_id": "test_dag",
|
|
"run_id": "prev_run",
|
|
"logical_date": "2024-01-14T12:00:00+00:00",
|
|
"start_date": "2024-01-14T12:05:00+00:00",
|
|
"end_date": "2024-01-14T12:10:00+00:00",
|
|
"state": "success",
|
|
"try_number": 1,
|
|
"map_index": -1,
|
|
"duration": 300.0,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_previous(
|
|
dag_id="test_dag", task_id="test_task", logical_date=logical_date, state="success"
|
|
)
|
|
|
|
assert result.task_instance.state == "success"
|
|
|
|
def test_get_previous_with_map_index_filter(self):
|
|
"""Test get_previous functionality with map_index filtering."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/task-instances/previous/test_dag/test_task":
|
|
assert request.url.params == httpx.QueryParams(
|
|
map_index="0",
|
|
)
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"task_id": "test_task",
|
|
"dag_id": "test_dag",
|
|
"run_id": "prev_run",
|
|
"logical_date": "2024-01-14T12:00:00+00:00",
|
|
"start_date": "2024-01-14T12:05:00+00:00",
|
|
"end_date": "2024-01-14T12:10:00+00:00",
|
|
"state": "success",
|
|
"try_number": 1,
|
|
"map_index": 0,
|
|
"duration": 300.0,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_previous(dag_id="test_dag", task_id="test_task", map_index=0)
|
|
|
|
assert result.task_instance.map_index == 0
|
|
|
|
def test_get_previous_not_found(self):
|
|
"""Test get_previous when no previous TI exists returns None."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/task-instances/previous/test_dag/test_task":
|
|
# Return None (null) when no previous TI found
|
|
return httpx.Response(status_code=200, content="null")
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_previous(dag_id="test_dag", task_id="test_task")
|
|
|
|
assert isinstance(result, PreviousTIResult)
|
|
assert result.task_instance is None
|
|
|
|
|
|
class TestVariableOperations:
|
|
"""
|
|
Test that the VariableOperations class works as expected. While the operations are simple, it
|
|
still catches the basic functionality of the client for variables including endpoint and
|
|
response parsing.
|
|
"""
|
|
|
|
def test_variable_get_success(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
# Simulate a successful response from the server with a variable
|
|
# ...including a validation that retry really works
|
|
call_count = 0
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count < 2:
|
|
return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
|
|
if request.url.path == "/variables/test_key":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"key": "test_key", "value": "test_value"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.variables.get(key="test_key")
|
|
|
|
assert isinstance(result, VariableResponse)
|
|
assert result.key == "test_key"
|
|
assert result.value == "test_value"
|
|
assert call_count == 2
|
|
|
|
def test_variable_not_found(self, cap_structlog):
|
|
# Simulate a 404 response from the server
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/variables/non_existent_var":
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={
|
|
"detail": {
|
|
"message": "Variable with key 'non_existent_var' not found",
|
|
"reason": "not_found",
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with cap_structlog.at_level("debug"):
|
|
resp = client.variables.get(key="non_existent_var")
|
|
|
|
assert isinstance(resp, ErrorResponse)
|
|
assert resp.error == ErrorType.VARIABLE_NOT_FOUND
|
|
assert resp.detail == {"key": "non_existent_var"}
|
|
# Verify the log is at debug level, not error (#52771)
|
|
assert {"log_level": "debug", "event": "Variable not found"} in cap_structlog
|
|
|
|
def test_variable_get_500_error(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
# Simulate a response from the server returning a 500 error
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/variables/test_key":
|
|
return httpx.Response(
|
|
status_code=500,
|
|
headers=[("content-Type", "application/json")],
|
|
json={
|
|
"reason": "internal_server_error",
|
|
"message": "Internal Server Error",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
with pytest.raises(ServerResponseError):
|
|
client.variables.get(
|
|
key="test_key",
|
|
)
|
|
|
|
@pytest.mark.parametrize("status_code", [401, 403])
|
|
def test_variable_get_authz_returns_permission_denied(self, status_code):
|
|
"""401/403 from the API server is reported as PERMISSION_DENIED, not raised.
|
|
|
|
The ExecutionAPISecretsBackend uses this distinction to refuse fallback
|
|
to a less-restrictive backend on an explicit deny.
|
|
"""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=status_code, json={"detail": "Forbidden"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
resp = client.variables.get(key="denied_var")
|
|
assert isinstance(resp, ErrorResponse)
|
|
assert resp.error == ErrorType.PERMISSION_DENIED
|
|
assert resp.detail == {"key": "denied_var", "status_code": status_code}
|
|
|
|
def test_variable_set_success(self):
|
|
# Simulate a successful response from the server when putting a variable
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/variables/test_key":
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"message": "Variable successfully set"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
result = client.variables.set(key="test_key", value="test_value", description="test_description")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_variable_delete_success(self):
|
|
# Simulate a successful response from the server when deleting a variable
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.method == "DELETE" and request.url.path == "/variables/test_key":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"count": 1},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
result = client.variables.delete(key="test_key")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
|
|
class TestXCOMOperations:
|
|
"""
|
|
Test that the XComOperations class works as expected. While the operations are simple, it
|
|
still catches the basic functionality of the client for xcoms including endpoint and
|
|
response parsing.
|
|
"""
|
|
|
|
@pytest.mark.parametrize(
|
|
"value",
|
|
[
|
|
pytest.param("value1", id="string-value"),
|
|
pytest.param({"key1": "value1"}, id="dict-value"),
|
|
pytest.param('{"key1": "value1"}', id="dict-str-value"),
|
|
pytest.param(["value1", "value2"], id="list-value"),
|
|
pytest.param({"key": "test_key", "value": {"key2": "value2"}}, id="nested-dict-value"),
|
|
],
|
|
)
|
|
def test_xcom_get_success(self, value):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
# Simulate a successful response from the server when getting an xcom
|
|
# ...including a validation that retry really works
|
|
call_count = 0
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count < 3:
|
|
return httpx.Response(status_code=500, json={"detail": "Internal Server Error"})
|
|
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"key": "test_key", "value": value},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.get(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
)
|
|
assert isinstance(result, XComResponse)
|
|
assert result.key == "test_key"
|
|
assert result.value == value
|
|
assert call_count == 3
|
|
|
|
def test_xcom_get_success_with_map_index(self):
|
|
# Simulate a successful response from the server when getting an xcom with map_index passed
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if (
|
|
request.url.path == "/xcoms/dag_id/run_id/task_id/key"
|
|
and request.url.params.get("map_index") == "2"
|
|
):
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"key": "test_key", "value": "test_value"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.get(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
map_index=2,
|
|
)
|
|
assert isinstance(result, XComResponse)
|
|
assert result.key == "test_key"
|
|
assert result.value == "test_value"
|
|
|
|
def test_xcom_get_success_with_include_prior_dates(self):
|
|
# Simulate a successful response from the server when getting an xcom with include_prior_dates passed
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/xcoms/dag_id/run_id/task_id/key" and request.url.params.get(
|
|
"include_prior_dates"
|
|
):
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"key": "test_key", "value": "test_value"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.get(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
include_prior_dates=True,
|
|
)
|
|
assert isinstance(result, XComResponse)
|
|
assert result.key == "test_key"
|
|
assert result.value == "test_value"
|
|
|
|
def test_xcom_get_500_error(self):
|
|
with time_machine.travel("2023-01-01T00:00:00Z", tick=False):
|
|
# Simulate a successful response from the server returning a 500 error
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
|
|
return httpx.Response(
|
|
status_code=500,
|
|
headers=[("content-Type", "application/json")],
|
|
json={
|
|
"reason": "invalid_format",
|
|
"message": "XCom value is not a valid JSON",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
with pytest.raises(ServerResponseError):
|
|
client.xcoms.get(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"values",
|
|
[
|
|
pytest.param("value1", id="string-value"),
|
|
pytest.param({"key1": "value1"}, id="dict-value"),
|
|
pytest.param(["value1", "value2"], id="list-value"),
|
|
pytest.param({"key": "test_key", "value": {"key2": "value2"}}, id="nested-dict-value"),
|
|
],
|
|
)
|
|
def test_xcom_set_success(self, values):
|
|
# Simulate a successful response from the server when setting an xcom
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
|
|
assert json.loads(request.read()) == values
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"message": "XCom successfully set"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.set(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
value=values,
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_xcom_set_with_map_index(self):
|
|
# Simulate a successful response from the server when setting an xcom with map_index passed
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if (
|
|
request.url.path == "/xcoms/dag_id/run_id/task_id/key"
|
|
and request.url.params.get("map_index") == "2"
|
|
):
|
|
assert json.loads(request.read()) == "value1"
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"message": "XCom successfully set"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.set(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
value="value1",
|
|
map_index=2,
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_xcom_set_with_mapped_length(self):
|
|
# Simulate a successful response from the server when setting an xcom with mapped_length
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if (
|
|
request.url.path == "/xcoms/dag_id/run_id/task_id/key"
|
|
and request.url.params.get("map_index") == "2"
|
|
and request.url.params.get("mapped_length") == "3"
|
|
):
|
|
assert json.loads(request.read()) == "value1"
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={"message": "XCom successfully set"},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.xcoms.set(
|
|
dag_id="dag_id",
|
|
run_id="run_id",
|
|
task_id="task_id",
|
|
key="key",
|
|
value="value1",
|
|
map_index=2,
|
|
mapped_length=3,
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
|
|
class TestConnectionOperations:
|
|
"""
|
|
Test that the TestConnectionOperations class works as expected. While the operations are simple, it
|
|
still catches the basic functionality of the client for connections including endpoint and
|
|
response parsing.
|
|
"""
|
|
|
|
def test_connection_get_success(self):
|
|
# Simulate a successful response from the server with a connection
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/connections/test_conn":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"conn_id": "test_conn",
|
|
"conn_type": "mysql",
|
|
"host": None,
|
|
"schema": None,
|
|
"login": None,
|
|
"password": None,
|
|
"port": None,
|
|
"extra": None,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.connections.get(conn_id="test_conn")
|
|
|
|
assert isinstance(result, ConnectionResponse)
|
|
assert result.conn_id == "test_conn"
|
|
assert result.conn_type == "mysql"
|
|
|
|
def test_connection_get_404_not_found(self):
|
|
# Simulate a successful response from the server with a connection
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/connections/test_conn":
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={
|
|
"reason": "not_found",
|
|
"message": "Connection with ID test_conn not found",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.connections.get(conn_id="test_conn")
|
|
|
|
assert isinstance(result, ErrorResponse)
|
|
assert result.error == ErrorType.CONNECTION_NOT_FOUND
|
|
|
|
def test_connection_get_url_encodes_conn_id(self):
|
|
requests_seen = []
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
requests_seen.append(request)
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"conn_id": "dev/my_conn",
|
|
"conn_type": "http",
|
|
"host": None,
|
|
"schema": None,
|
|
"login": None,
|
|
"password": None,
|
|
"port": None,
|
|
"extra": None,
|
|
},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.connections.get(conn_id="dev/my_conn")
|
|
|
|
assert isinstance(result, ConnectionResponse)
|
|
assert requests_seen[0].url.raw_path == b"/connections/dev%2Fmy_conn"
|
|
|
|
@pytest.mark.parametrize("status_code", [401, 403])
|
|
def test_connection_get_authz_returns_permission_denied(self, status_code):
|
|
"""401/403 from the API server is reported as PERMISSION_DENIED, not raised."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=status_code, json={"detail": "Forbidden"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.connections.get(conn_id="denied_conn")
|
|
assert isinstance(result, ErrorResponse)
|
|
assert result.error == ErrorType.PERMISSION_DENIED
|
|
assert result.detail == {"conn_id": "denied_conn", "status_code": status_code}
|
|
|
|
|
|
class TestAssetEventOperations:
|
|
@pytest.mark.parametrize(
|
|
"request_params",
|
|
[
|
|
({"name": "this_asset", "uri": "s3://bucket/key"}),
|
|
({"alias_name": "this_asset_alias"}),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("created_dagruns", "expected_created_dagruns"),
|
|
[
|
|
pytest.param([], 0, id="without-created-dagrun"),
|
|
pytest.param(
|
|
[
|
|
{
|
|
"dag_id": "created_dag",
|
|
"run_id": "queued_run",
|
|
"logical_date": "2023-01-01T00:00:00Z",
|
|
"start_date": None,
|
|
"end_date": None,
|
|
"state": "queued",
|
|
"data_interval_start": None,
|
|
"data_interval_end": None,
|
|
"partition_key": None,
|
|
}
|
|
],
|
|
1,
|
|
id="queued-created-dagrun-without-start-date",
|
|
),
|
|
],
|
|
)
|
|
def test_by_name_get_success(self, request_params, created_dagruns, expected_created_dagruns):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
if request.url.path == "/asset-events/by-asset":
|
|
assert params.get("name") == request_params.get("name")
|
|
assert params.get("uri") == request_params.get("uri")
|
|
elif request.url.path == "/asset-events/by-asset-alias":
|
|
assert params.get("name") == request_params.get("alias_name")
|
|
else:
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"asset_events": [
|
|
{
|
|
"id": 1,
|
|
"asset": {
|
|
"name": "this_asset",
|
|
"uri": "s3://bucket/key",
|
|
"group": "asset",
|
|
},
|
|
"created_dagruns": created_dagruns,
|
|
"timestamp": "2023-01-01T00:00:00Z",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(**request_params)
|
|
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 1
|
|
assert result.asset_events[0].asset.name == "this_asset"
|
|
assert result.asset_events[0].asset.uri == "s3://bucket/key"
|
|
assert len(result.asset_events[0].created_dagruns) == expected_created_dagruns
|
|
if expected_created_dagruns:
|
|
assert result.asset_events[0].created_dagruns[0].start_date is None
|
|
|
|
def test_partition_key_exact_match_param_passed(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
assert params.get("partition_key") == "2024-01-01"
|
|
assert "partition_key_regexp_pattern" not in params
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"asset_events": [
|
|
{
|
|
"id": 1,
|
|
"asset": {
|
|
"name": "this_asset",
|
|
"uri": "s3://bucket/key",
|
|
"group": "asset",
|
|
},
|
|
"created_dagruns": [],
|
|
"timestamp": "2023-01-01T00:00:00Z",
|
|
"partition_key": "2024-01-01",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = make_client(httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(name="this_asset", partition_key="2024-01-01")
|
|
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 1
|
|
|
|
def test_partition_key_regexp_pattern_param_passed(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
assert params.get("partition_key_regexp_pattern") == "^2024-01-"
|
|
assert "partition_key" not in params
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"asset_events": []},
|
|
)
|
|
|
|
client = make_client(httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(name="this_asset", partition_key_regexp_pattern="^2024-01-")
|
|
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 0
|
|
|
|
def test_partition_key_regexp_pattern_with_other_params(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
assert params.get("partition_key_regexp_pattern") == r"^us\|2024-.*"
|
|
assert "partition_key" not in params
|
|
assert params.get("after") == "2023-06-01T00:00:00+00:00"
|
|
assert params.get("limit") == "5"
|
|
assert params.get("ascending") == "false"
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"asset_events": []},
|
|
)
|
|
|
|
client = make_client(httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(
|
|
name="this_asset",
|
|
partition_key_regexp_pattern=r"^us\|2024-.*",
|
|
after=datetime(2023, 6, 1, tzinfo=UTC),
|
|
limit=5,
|
|
ascending=False,
|
|
)
|
|
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 0
|
|
|
|
def test_partition_key_with_alias(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/asset-events/by-asset-alias"
|
|
params = request.url.params
|
|
assert params.get("name") == "my_alias"
|
|
assert params.get("partition_key") == "us-east|2024-01-01"
|
|
assert "partition_key_regexp_pattern" not in params
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"asset_events": []},
|
|
)
|
|
|
|
client = make_client(httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(alias_name="my_alias", partition_key="us-east|2024-01-01")
|
|
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 0
|
|
|
|
def test_extra_dict_param_passed(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
extra_pairs = [v for k, v in params.multi_items() if k == "extra"]
|
|
assert sorted(extra_pairs) == sorted(["region=us", "env=prod"])
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"asset_events": [
|
|
{
|
|
"id": 1,
|
|
"asset": {"name": "a", "uri": "s3://b", "group": "asset"},
|
|
"created_dagruns": [],
|
|
"timestamp": "2023-01-01T00:00:00Z",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(name="a", uri="s3://b", extra={"region": "us", "env": "prod"})
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 1
|
|
|
|
def test_extra_dict_with_alias(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
assert request.url.path == "/asset-events/by-asset-alias"
|
|
extra_pairs = [v for k, v in params.multi_items() if k == "extra"]
|
|
assert extra_pairs == ["env=prod"]
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"asset_events": [
|
|
{
|
|
"id": 1,
|
|
"asset": {"name": "a", "uri": "s3://b", "group": "asset"},
|
|
"created_dagruns": [],
|
|
"timestamp": "2023-01-01T00:00:00Z",
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(alias_name="my_alias", extra={"env": "prod"})
|
|
assert isinstance(result, AssetEventsResponse)
|
|
assert len(result.asset_events) == 1
|
|
|
|
def test_extra_none_not_sent(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
params = request.url.params
|
|
assert "extra" not in params
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"asset_events": []},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_events.get(name="a", uri="s3://b")
|
|
assert isinstance(result, AssetEventsResponse)
|
|
|
|
|
|
class TestAssetOperations:
|
|
@pytest.mark.parametrize(
|
|
"request_params",
|
|
[
|
|
({"name": "this_asset"}),
|
|
({"uri": "s3://bucket/key"}),
|
|
],
|
|
)
|
|
def test_by_name_get_success(self, request_params):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path in ("/assets/by-name", "/assets/by-uri"):
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"name": "this_asset",
|
|
"uri": "s3://bucket/key",
|
|
"group": "asset",
|
|
"extra": {"foo": "bar"},
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.assets.get(**request_params)
|
|
|
|
assert isinstance(result, AssetResponse)
|
|
assert result.name == "this_asset"
|
|
assert result.uri == "s3://bucket/key"
|
|
|
|
@pytest.mark.parametrize(
|
|
"request_params",
|
|
[
|
|
({"name": "this_asset"}),
|
|
({"uri": "s3://bucket/key"}),
|
|
],
|
|
)
|
|
def test_by_name_get_404_not_found(self, request_params):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path in ("/assets/by-name", "/assets/by-uri"):
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={
|
|
"detail": {
|
|
"message": "Asset with name non_existent not found",
|
|
"reason": "not_found",
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.assets.get(**request_params)
|
|
|
|
assert isinstance(result, ErrorResponse)
|
|
assert result.error == ErrorType.ASSET_NOT_FOUND
|
|
|
|
def test_get_by_alias_returns_list(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/assets/by-alias"
|
|
assert request.url.params["alias_name"] == "my_alias"
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json=[
|
|
{"name": "asset_a", "uri": "s3://bucket/a", "group": "asset", "extra": {}},
|
|
{"name": "asset_b", "uri": "s3://bucket/b", "group": "asset", "extra": {}},
|
|
],
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.assets.get_by_alias("my_alias")
|
|
|
|
assert isinstance(result, AssetsByAliasResult)
|
|
assert len(result.assets) == 2
|
|
assert isinstance(result.assets[0], AssetResponse)
|
|
assert isinstance(result.assets[1], AssetResponse)
|
|
assert result.assets[0].name == "asset_a"
|
|
assert result.assets[1].name == "asset_b"
|
|
|
|
def test_get_by_alias_returns_empty_for_unknown_alias(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(status_code=200, json=[])
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.assets.get_by_alias("unknown_alias")
|
|
|
|
assert result.assets == []
|
|
|
|
|
|
class TestDagRunOperations:
|
|
def test_trigger(self):
|
|
# Simulate a successful response from the server when triggering a dag run
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/test_trigger/test_run_id":
|
|
actual_body = json.loads(request.read())
|
|
assert actual_body["logical_date"] == "2025-01-01T00:00:00Z"
|
|
assert actual_body["conf"] == {}
|
|
|
|
# Since the value for `reset_dag_run` is default, it should not be present in the request body
|
|
assert "reset_dag_run" not in actual_body
|
|
return httpx.Response(status_code=204)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.trigger(
|
|
dag_id="test_trigger",
|
|
run_id="test_run_id",
|
|
logical_date=timezone.datetime(2025, 1, 1),
|
|
run_after=timezone.datetime(2025, 1, 2),
|
|
)
|
|
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_trigger_conflict(self):
|
|
"""Test that if the dag run already exists, the client returns an error when default reset_dag_run=False"""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/test_trigger_conflict/test_run_id":
|
|
return httpx.Response(
|
|
status_code=409,
|
|
json={
|
|
"detail": {
|
|
"reason": "already_exists",
|
|
"message": "A Dag Run already exists for Dag test_trigger_conflict with run id test_run_id",
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.trigger(dag_id="test_trigger_conflict", run_id="test_run_id")
|
|
|
|
assert result == ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS)
|
|
|
|
def test_trigger_conflict_reset_dag_run(self):
|
|
"""Test that if dag run already exists and reset_dag_run=True, the client clears the dag run"""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id":
|
|
return httpx.Response(
|
|
status_code=409,
|
|
json={
|
|
"detail": {
|
|
"reason": "already_exists",
|
|
"message": "A Dag Run already exists for Dag test_trigger_conflict with run id test_run_id",
|
|
}
|
|
},
|
|
)
|
|
if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id/clear":
|
|
return httpx.Response(status_code=204)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.trigger(
|
|
dag_id="test_trigger_conflict_reset",
|
|
run_id="test_run_id",
|
|
reset_dag_run=True,
|
|
)
|
|
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_clear(self):
|
|
"""Test that the client can clear a dag run"""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/test_clear/test_run_id/clear":
|
|
return httpx.Response(status_code=204)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.clear(dag_id="test_clear", run_id="test_run_id")
|
|
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_get_state(self):
|
|
"""Test that the client can get the state of a dag run"""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/test_state/test_run_id/state":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"state": "running"},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_state(dag_id="test_state", run_id="test_run_id")
|
|
|
|
assert result == DagRunStateResponse(state=DagRunState.RUNNING)
|
|
|
|
def test_get_count_basic(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/count":
|
|
assert request.url.params["dag_id"] == "test_dag"
|
|
return httpx.Response(status_code=200, json=1)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_count(dag_id="test_dag")
|
|
assert result.count == 1
|
|
|
|
def test_get_count_with_states(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/count":
|
|
assert request.url.params["dag_id"] == "test_dag"
|
|
assert request.url.params.get_list("states") == ["success", "failed"]
|
|
return httpx.Response(status_code=200, json=2)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_count(dag_id="test_dag", states=["success", "failed"])
|
|
assert result.count == 2
|
|
|
|
def test_get_count_with_logical_dates(self):
|
|
logical_dates = [timezone.datetime(2025, 1, 1), timezone.datetime(2025, 1, 2)]
|
|
logical_dates_str = [d.isoformat() for d in logical_dates]
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/count":
|
|
assert request.url.params["dag_id"] == "test_dag"
|
|
assert request.url.params.get_list("logical_dates") == logical_dates_str
|
|
return httpx.Response(status_code=200, json=2)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_count(
|
|
dag_id="test_dag", logical_dates=[timezone.datetime(2025, 1, 1), timezone.datetime(2025, 1, 2)]
|
|
)
|
|
assert result.count == 2
|
|
|
|
def test_get_count_with_run_ids(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/count":
|
|
assert request.url.params["dag_id"] == "test_dag"
|
|
assert request.url.params.get_list("run_ids") == ["run1", "run2"]
|
|
return httpx.Response(status_code=200, json=2)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_count(dag_id="test_dag", run_ids=["run1", "run2"])
|
|
assert result.count == 2
|
|
|
|
def test_get_previous_basic(self):
|
|
"""Test basic get_previous functionality with dag_id and logical_date."""
|
|
logical_date = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/previous":
|
|
assert request.url.params == httpx.QueryParams(
|
|
dag_id="test_dag",
|
|
logical_date=logical_date.isoformat(),
|
|
)
|
|
# Return complete DagRun data
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"dag_id": "test_dag",
|
|
"run_id": "prev_run",
|
|
"logical_date": "2024-01-14T12:00:00+00:00",
|
|
"start_date": "2024-01-14T12:05:00+00:00",
|
|
"run_after": "2024-01-14T12:00:00+00:00",
|
|
"run_type": "scheduled",
|
|
"state": "success",
|
|
"consumed_asset_events": [],
|
|
"data_interval_start": None,
|
|
"data_interval_end": None,
|
|
"end_date": None,
|
|
"partition_key": None,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_previous(dag_id="test_dag", logical_date=logical_date)
|
|
|
|
assert isinstance(result, PreviousDagRunResult)
|
|
assert result.dag_run.dag_id == "test_dag"
|
|
assert result.dag_run.run_id == "prev_run"
|
|
assert result.dag_run.state == "success"
|
|
|
|
def test_get_previous_with_state_filter(self):
|
|
"""Test get_previous functionality with state filtering."""
|
|
logical_date = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/previous":
|
|
assert request.url.params == httpx.QueryParams(
|
|
dag_id="test_dag",
|
|
logical_date=logical_date.isoformat(),
|
|
state="success",
|
|
)
|
|
# Return complete DagRun data
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"dag_id": "test_dag",
|
|
"run_id": "prev_success_run",
|
|
"logical_date": "2024-01-14T12:00:00+00:00",
|
|
"start_date": "2024-01-14T12:05:00+00:00",
|
|
"run_after": "2024-01-14T12:00:00+00:00",
|
|
"run_type": "scheduled",
|
|
"state": "success",
|
|
"consumed_asset_events": [],
|
|
"data_interval_start": None,
|
|
"data_interval_end": None,
|
|
"end_date": None,
|
|
"partition_key": None,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_previous(dag_id="test_dag", logical_date=logical_date, state="success")
|
|
|
|
assert isinstance(result, PreviousDagRunResult)
|
|
assert result.dag_run.dag_id == "test_dag"
|
|
assert result.dag_run.run_id == "prev_success_run"
|
|
assert result.dag_run.state == "success"
|
|
|
|
def test_get_previous_not_found(self):
|
|
"""Test get_previous when no previous Dag run exists returns None."""
|
|
logical_date = datetime(2024, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dag-runs/previous":
|
|
assert request.url.params == httpx.QueryParams(
|
|
dag_id="test_dag",
|
|
logical_date=logical_date.isoformat(),
|
|
)
|
|
# Return None (null) when no previous Dag run found
|
|
return httpx.Response(status_code=200, content="null")
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dag_runs.get_previous(dag_id="test_dag", logical_date=logical_date)
|
|
|
|
assert isinstance(result, PreviousDagRunResult)
|
|
assert result.dag_run is None
|
|
|
|
|
|
class TestTaskRescheduleOperations:
|
|
def test_get_start_date(self):
|
|
"""Test that the client can get the start date of a task reschedule"""
|
|
ti_id = uuid6.uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/task-reschedules/{ti_id}/start_date":
|
|
assert request.url.params.get("try_number") == "1"
|
|
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json="2024-01-01T00:00:00Z",
|
|
)
|
|
return httpx.Response(status_code=422)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_instances.get_reschedule_start_date(id=ti_id, try_number=1)
|
|
|
|
assert isinstance(result, TaskRescheduleStartDate)
|
|
assert result.start_date == datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
|
|
|
|
|
|
class TestHITLOperations:
|
|
def test_add_response(self) -> None:
|
|
ti_id = uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path in (f"/hitlDetails/{ti_id}"):
|
|
return httpx.Response(
|
|
status_code=201,
|
|
json={
|
|
"ti_id": str(ti_id),
|
|
"options": ["Approval", "Reject"],
|
|
"subject": "This is subject",
|
|
"body": "This is body",
|
|
"defaults": ["Approval"],
|
|
"params": None,
|
|
"multiple": False,
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.hitl.add_response(
|
|
ti_id=ti_id,
|
|
options=["Approval", "Reject"],
|
|
subject="This is subject",
|
|
body="This is body",
|
|
defaults=["Approval"],
|
|
params=None,
|
|
multiple=False,
|
|
)
|
|
assert isinstance(result, HITLDetailRequest)
|
|
assert result.ti_id == ti_id
|
|
assert result.options == ["Approval", "Reject"]
|
|
assert result.subject == "This is subject"
|
|
assert result.body == "This is body"
|
|
assert result.defaults == ["Approval"]
|
|
assert result.params is None
|
|
assert result.multiple is False
|
|
assert result.assigned_users is None
|
|
|
|
def test_update_response(self, time_machine: TimeMachineFixture) -> None:
|
|
time_machine.move_to(datetime(2025, 7, 3, 0, 0, 0))
|
|
ti_id = uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path in (f"/hitlDetails/{ti_id}"):
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"chosen_options": ["Approval"],
|
|
"params_input": {},
|
|
"responded_by_user": {"id": "admin", "name": "admin"},
|
|
"response_received": True,
|
|
"responded_at": "2025-07-03T00:00:00Z",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.hitl.update_response(
|
|
ti_id=ti_id,
|
|
chosen_options=["Approve"],
|
|
params_input={},
|
|
)
|
|
assert isinstance(result, HITLDetailResponse)
|
|
assert result.response_received is True
|
|
assert result.chosen_options == ["Approval"]
|
|
assert result.params_input == {}
|
|
assert result.responded_by_user == HITLUser(id="admin", name="admin")
|
|
assert result.responded_at == timezone.datetime(2025, 7, 3, 0, 0, 0)
|
|
|
|
def test_get_detail_response(self, time_machine: TimeMachineFixture) -> None:
|
|
time_machine.move_to(datetime(2025, 7, 3, 0, 0, 0))
|
|
ti_id = uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path in (f"/hitlDetails/{ti_id}"):
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"chosen_options": ["Approval"],
|
|
"params_input": {},
|
|
"responded_by_user": {"id": "admin", "name": "admin"},
|
|
"response_received": True,
|
|
"responded_at": "2025-07-03T00:00:00Z",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=400, json={"detail": "Bad Request"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.hitl.get_detail_response(ti_id=ti_id)
|
|
assert isinstance(result, HITLDetailResponse)
|
|
assert result.response_received is True
|
|
assert result.chosen_options == ["Approval"]
|
|
assert result.params_input == {}
|
|
assert result.responded_by_user == HITLUser(id="admin", name="admin")
|
|
assert result.responded_at == timezone.datetime(2025, 7, 3, 0, 0, 0)
|
|
|
|
|
|
class TestSSLContextCaching:
|
|
@pytest.fixture(autouse=True)
|
|
def clear_ssl_context_cache(self):
|
|
Client._get_ssl_context_cached.cache_clear()
|
|
yield
|
|
Client._get_ssl_context_cached.cache_clear()
|
|
|
|
def test_cache_hit_on_same_parameters(self):
|
|
ca_file = certifi.where()
|
|
ctx1 = Client._get_ssl_context_cached(ca_file, None)
|
|
ctx2 = Client._get_ssl_context_cached(ca_file, None)
|
|
assert ctx1 is ctx2
|
|
|
|
def test_cache_miss_if_cache_cleared(self):
|
|
ca_file = certifi.where()
|
|
ctx1 = Client._get_ssl_context_cached(ca_file, None)
|
|
Client._get_ssl_context_cached.cache_clear()
|
|
ctx2 = Client._get_ssl_context_cached(ca_file, None)
|
|
assert ctx1 is not ctx2
|
|
|
|
def test_cache_miss_on_different_parameters(self):
|
|
ca_file = certifi.where()
|
|
|
|
ctx1 = Client._get_ssl_context_cached(ca_file, None)
|
|
ctx2 = Client._get_ssl_context_cached(ca_file, ca_file)
|
|
|
|
info = Client._get_ssl_context_cached.cache_info()
|
|
|
|
assert ctx1 is not ctx2
|
|
assert info.misses == 2
|
|
assert info.currsize == 2
|
|
|
|
|
|
class TestDagsOperations:
|
|
def test_get(self):
|
|
"""Test that the client can get a dag."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dags/test_dag":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={
|
|
"dag_id": "test_dag",
|
|
"is_paused": False,
|
|
"bundle_name": "dags-folder",
|
|
"bundle_version": "bundle-version",
|
|
"relative_fileloc": "dags/example.py",
|
|
"owners": "owner_1",
|
|
"tags": ["a_tag", "z_tag"],
|
|
"next_dagrun": "2026-04-13T00:00:00Z",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=200)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.dags.get(dag_id="test_dag")
|
|
|
|
assert result == DagResponse(
|
|
dag_id="test_dag",
|
|
is_paused=False,
|
|
bundle_name="dags-folder",
|
|
bundle_version="bundle-version",
|
|
relative_fileloc="dags/example.py",
|
|
owners="owner_1",
|
|
tags=["a_tag", "z_tag"],
|
|
next_dagrun=datetime(2026, 4, 13, tzinfo=UTC),
|
|
)
|
|
|
|
def test_get_url_quotes_dag_id_as_single_path_segment(self):
|
|
"""Test that Dag IDs cannot escape the dags API path."""
|
|
requests_seen = []
|
|
crafted_dag_id = "x/../../variables/secret_key"
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
requests_seen.append(request)
|
|
if request.url.path == "/variables/secret_key":
|
|
return httpx.Response(
|
|
status_code=200,
|
|
json={"key": "secret_key", "value": "super-secret-value"},
|
|
)
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={
|
|
"detail": {
|
|
"message": "The Dag was not found",
|
|
"reason": "not_found",
|
|
}
|
|
},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with pytest.raises(ServerResponseError) as exc_info:
|
|
client.dags.get(dag_id=crafted_dag_id)
|
|
|
|
assert requests_seen[0].url.raw_path == b"/dags/x%2F..%2F..%2Fvariables%2Fsecret_key"
|
|
assert requests_seen[0].url.path != "/variables/secret_key"
|
|
assert "super-secret-value" not in str(exc_info.value)
|
|
|
|
def test_get_not_found(self):
|
|
"""Test that getting a missing dag raises a server response error."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dags/missing_dag":
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={
|
|
"detail": {
|
|
"message": "The Dag with dag_id: `missing_dag` was not found",
|
|
"reason": "not_found",
|
|
}
|
|
},
|
|
)
|
|
return httpx.Response(status_code=200)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with pytest.raises(ServerResponseError) as exc_info:
|
|
client.dags.get(dag_id="missing_dag")
|
|
|
|
assert exc_info.value.response.status_code == 404
|
|
assert exc_info.value.detail == {
|
|
"message": "The Dag with dag_id: `missing_dag` was not found",
|
|
"reason": "not_found",
|
|
}
|
|
|
|
def test_get_server_error(self):
|
|
"""Test that a server error while getting a dag."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == "/dags/test_dag":
|
|
return httpx.Response(
|
|
status_code=500,
|
|
headers=[("content-Type", "application/json")],
|
|
json={
|
|
"reason": "internal_server_error",
|
|
"message": "Internal Server Error",
|
|
},
|
|
)
|
|
return httpx.Response(status_code=200)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
|
|
with pytest.raises(ServerResponseError):
|
|
client.dags.get(dag_id="test_dag")
|
|
|
|
|
|
class TestTaskStateOperations:
|
|
TI_ID = uuid7()
|
|
|
|
def test_get_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if request.url.path == f"/store/ti/{self.TI_ID}/job_id":
|
|
return httpx.Response(status_code=200, json={"value": "spark_app_001"})
|
|
return httpx.Response(status_code=400)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.get(ti_id=self.TI_ID, key="job_id")
|
|
|
|
assert isinstance(result, TaskStateStoreResponse)
|
|
assert result.value == "spark_app_001"
|
|
|
|
def test_get_url_encodes_key_with_slash(self):
|
|
requests_seen = []
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
requests_seen.append(request)
|
|
return httpx.Response(status_code=200, json={"value": "spark_app_001"})
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_state_store.get(ti_id=self.TI_ID, key="spark/job_id")
|
|
|
|
assert b"%2F" in requests_seen[0].url.raw_path
|
|
assert requests_seen[0].url.raw_path == f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
|
|
|
|
def test_get_returns_error_response_on_404(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={"detail": {"reason": "not_found", "message": "Task state key 'job_id' not found"}},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.get(ti_id=self.TI_ID, key="job_id")
|
|
assert isinstance(result, ErrorResponse)
|
|
assert result.error == ErrorType.TASK_STORE_NOT_FOUND
|
|
|
|
def test_set_success(self):
|
|
expires = datetime(2026, 6, 13, 12, 0, 0, tzinfo=UTC)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "PUT"
|
|
assert request.url.path == f"/store/ti/{self.TI_ID}/job_id"
|
|
assert b'"value":"spark_app_001"' in request.content
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.set(
|
|
ti_id=self.TI_ID, key="job_id", value="spark_app_001", expires_at=expires
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_set_url_encodes_key_with_slash(self):
|
|
requests_seen = []
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
requests_seen.append(request)
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_state_store.set(
|
|
ti_id=self.TI_ID, key="spark/job_id", value="spark_app_001", expires_at=None
|
|
)
|
|
|
|
assert b"%2F" in requests_seen[0].url.raw_path
|
|
assert requests_seen[0].url.raw_path == f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
|
|
|
|
def test_set_with_expires_at_sends_field(self):
|
|
"""expires_at is forwarded as an ISO datetime string in the request body."""
|
|
expires = datetime(2026, 5, 21, 12, 0, 0, tzinfo=UTC)
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
body = json.loads(request.content)
|
|
assert body["value"] == "spark_app_001"
|
|
assert body["expires_at"] == "2026-05-21T12:00:00Z"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.set(
|
|
ti_id=self.TI_ID, key="job_id", value="spark_app_001", expires_at=expires
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_set_with_never_expire_sends_null_expires_at(self):
|
|
"""NEVER_EXPIRE sends expires_at=null — stored as NULL in DB, GC skips it."""
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
body = json.loads(request.content)
|
|
assert body.get("expires_at") is None
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.set(
|
|
ti_id=self.TI_ID,
|
|
key="job_id",
|
|
value="v",
|
|
expires_at=None,
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_delete_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == f"/store/ti/{self.TI_ID}/job_id"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.delete(ti_id=self.TI_ID, key="job_id")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_delete_url_encodes_key_with_slash(self):
|
|
requests_seen = []
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
requests_seen.append(request)
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.task_state_store.delete(ti_id=self.TI_ID, key="spark/job_id")
|
|
|
|
assert b"%2F" in requests_seen[0].url.raw_path
|
|
assert requests_seen[0].url.raw_path == f"/store/ti/{self.TI_ID}/spark%2Fjob_id".encode()
|
|
|
|
def test_clear_sends_delete_request(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == f"/store/ti/{self.TI_ID}"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.task_state_store.clear(ti_id=self.TI_ID)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
|
|
class TestAssetStateOperations:
|
|
def test_get_by_name_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if (
|
|
request.url.path == "/store/asset/by-name/value"
|
|
and request.url.params["name"] == "test_asset"
|
|
and request.url.params["key"] == "watermark"
|
|
):
|
|
return httpx.Response(status_code=200, json={"value": "2026-04-30T00:00:00Z"})
|
|
return httpx.Response(status_code=400)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.get(key="watermark", name="test_asset")
|
|
|
|
assert isinstance(result, AssetStateStoreResponse)
|
|
assert result.value == "2026-04-30T00:00:00Z"
|
|
|
|
def test_get_by_uri_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
if (
|
|
request.url.path == "/store/asset/by-uri/value"
|
|
and request.url.params["uri"] == "s3://bucket/key"
|
|
and request.url.params["key"] == "watermark"
|
|
):
|
|
return httpx.Response(status_code=200, json={"value": "2026-04-30T00:00:00Z"})
|
|
return httpx.Response(status_code=400)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.get(key="watermark", uri="s3://bucket/key")
|
|
|
|
assert isinstance(result, AssetStateStoreResponse)
|
|
assert result.value == "2026-04-30T00:00:00Z"
|
|
|
|
def test_get_returns_error_response_on_404(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code=404,
|
|
json={"detail": {"reason": "not_found", "message": "Asset state key 'watermark' not found"}},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.get(key="watermark", name="test_asset")
|
|
assert isinstance(result, ErrorResponse)
|
|
assert result.error == ErrorType.ASSET_STORE_NOT_FOUND
|
|
|
|
def test_set_by_name_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "PUT"
|
|
assert request.url.path == "/store/asset/by-name/value"
|
|
assert request.url.params["name"] == "test_asset"
|
|
assert request.url.params["key"] == "watermark"
|
|
assert b'"value":"2026-04-30T00:00:00Z"' in request.content
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.set(
|
|
key="watermark", value="2026-04-30T00:00:00Z", name="test_asset"
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_set_by_uri_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "PUT"
|
|
assert request.url.path == "/store/asset/by-uri/value"
|
|
assert request.url.params["uri"] == "s3://bucket/key"
|
|
assert request.url.params["key"] == "watermark"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.set(
|
|
key="watermark", value="2026-04-30T00:00:00Z", uri="s3://bucket/key"
|
|
)
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_delete_by_name_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == "/store/asset/by-name/value"
|
|
assert request.url.params["name"] == "test_asset"
|
|
assert request.url.params["key"] == "watermark"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.delete(key="watermark", name="test_asset")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_delete_by_uri_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == "/store/asset/by-uri/value"
|
|
assert request.url.params["uri"] == "s3://bucket/key"
|
|
assert request.url.params["key"] == "watermark"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.delete(key="watermark", uri="s3://bucket/key")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_clear_by_name_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == "/store/asset/by-name/clear"
|
|
assert request.url.params["name"] == "test_asset"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.clear(name="test_asset")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
def test_clear_by_uri_success(self):
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "DELETE"
|
|
assert request.url.path == "/store/asset/by-uri/clear"
|
|
assert request.url.params["uri"] == "s3://bucket/key"
|
|
return httpx.Response(status_code=204)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
result = client.asset_state_store.clear(uri="s3://bucket/key")
|
|
assert result == OKResponse(ok=True)
|
|
|
|
|
|
class TestCallbackOperations:
|
|
def test_run_exchanges_token(self):
|
|
callback_id = uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
assert request.method == "PATCH"
|
|
assert request.url.path == f"/callbacks/{callback_id}/run"
|
|
return httpx.Response(
|
|
status_code=204,
|
|
headers={"Refreshed-API-Token": "execution-token"},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
client.callbacks.run(callback_id)
|
|
|
|
assert client.auth is not None
|
|
assert client.auth.token == "execution-token"
|
|
|
|
def test_run_raises_on_conflict(self):
|
|
callback_id = uuid7()
|
|
|
|
def handle_request(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
status_code=409,
|
|
json={"detail": {"reason": "invalid_state", "previous_state": "running"}},
|
|
)
|
|
|
|
client = make_client(transport=httpx.MockTransport(handle_request))
|
|
with pytest.raises(ServerResponseError) as err:
|
|
client.callbacks.run(callback_id)
|
|
assert err.value.response.status_code == 409
|