Google Git Google Git
Sign in
apache/airflow/refs/heads/hardened-python-base-images/./task-sdk/tests/task_sdk/api/test_client.py
blob: 62fde3373e007eb4d74707cb651f5c972b47b4b3 [file]
# 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 httpx2
import pytest
import time_machine
import truststore
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
@pytest.fixture
def ca_file(tmp_path):
from datetime import timedelta
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID
key = ec.generate_private_key(ec.SECP256R1())
name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "test-ca")])
now = datetime.now(UTC)
cert = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(name)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now)
.not_valid_after(now + timedelta(days=1))
.add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
.sign(key, hashes.SHA256())
)
path = tmp_path / "ca.pem"
path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
return str(path)
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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
with pytest.raises(FileNotFoundError) as err:
make_client(httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
with pytest.raises(FileNotFoundError) as err:
make_client(httpx2.MockTransport(handle_request))
assert isinstance(err.value, FileNotFoundError)
@mock.patch("airflow.sdk.api.client.API_CLIENT_USE_PUBLIC_CERTS", True)
def test_use_public_certs(self, ca_file):
ctx = Client._get_ssl_context_cached(ca_file, None)
assert isinstance(ctx, truststore.SSLContext)
@mock.patch("airflow.sdk.api.client.API_TIMEOUT", 60.0)
def test_timeout_configuration(self):
def handle_request(request: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
client = make_client(httpx2.MockTransport(handle_request))
assert client.timeout == httpx2.Timeout(60.0)
def test_timeout_can_be_overridden(self):
def handle_request(request: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
client = Client(
base_url="test://server", token="", transport=httpx2.MockTransport(handle_request), timeout=120.0
)
assert client.timeout == httpx2.Timeout(120.0)
def test_error_parsing(self):
responses = [
httpx2.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 = [httpx2.Response(422, content=b"Internal Server Error")]
client = make_client_w_responses(responses)
with pytest.raises(httpx2.HTTPStatusError) as err:
client.get("http://error")
assert not isinstance(err.value, ServerResponseError)
def test_error_parsing_other_json(self):
responses = [httpx2.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 = [httpx2.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 = [httpx2.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[httpx2.Response] = [
*[httpx2.Response(500, text="Internal Server Error")] * 6,
httpx2.Response(200, json={"detail": "Recovered from error - but will fail before"}),
httpx2.Response(400, json={"detail": "Should not get here"}),
]
client = make_client_w_responses(responses)
with pytest.raises(httpx2.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[httpx2.Response] = [
*[httpx2.Response(500, text="Internal Server Error")] * 2,
httpx2.Response(200, json={"detail": "Recovered from error"}),
httpx2.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[httpx2.Response] = [
httpx2.Response(422, json={"detail": "Somehow this is a bad request"}),
httpx2.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[httpx2.Response] = [
httpx2.Response(200, json={"detail": "Recovered from error"}),
httpx2.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[httpx2.Response] = [
httpx2.Response(200, json={"ok": "1"}),
httpx2.Response(404, json={"var": "not_found"}, headers={"Refreshed-API-Token": "abc"}),
httpx2.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 = httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
captured: dict[str, object] = {}
real_init = httpx2.Client.__init__
def spy_init(self, *args, **kwargs):
captured["cert"] = kwargs.get("cert")
return real_init(self, *args, **kwargs)
with mock.patch.object(httpx2.Client, "__init__", spy_init):
make_client(httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200)
with pytest.raises(ValueError, match="Both client_ssl_cert and client_ssl_key must be set"):
make_client(httpx2.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: httpx2.Request) -> httpx2.Response:
nonlocal call_count
call_count += 1
if call_count < 3:
return httpx2.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 httpx2.Response(
status_code=200,
json=ti_context.model_dump(mode="json"),
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-instances/{ti_id}/run":
return httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-instances/{ti_id}/run":
return httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=204,
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-instances/{ti_id}/heartbeat":
actual_body = json.loads(request.read())
assert actual_body["pid"] == 100
return httpx2.Response(
status_code=204,
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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[httpx2.Response] = [
httpx2.Response(500, text="Internal Server Error"),
httpx2.Response(204),
]
client = make_client_w_responses(responses)
with pytest.raises(httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-instances/{ti_id}/dag-run-note":
assert json.loads(request.read()) == {"note": "Updated from task runtime"}
return httpx2.Response(status_code=204)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=204,
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=204,
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=204,
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-instances/{TI_ID}/rtif":
return httpx2.Response(
status_code=201,
json={"message": "Rendered task instance fields successfully set"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.url.path == "/task-instances/count"
assert request.url.params.get("dag_id") == "test_dag"
return httpx2.Response(200, json=5)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(200, json=10)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
200, json={"task_states": {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}}
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
200, json={"task_states": {"run_id": {"group1.task1": "success", "group1.task2": "failed"}}}
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/task-instances/previous/test_dag/test_task":
assert request.url.params == httpx2.QueryParams(
logical_date=logical_date.isoformat(),
map_index="-1",
)
# Return complete TI data
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/task-instances/previous/test_dag/test_task":
assert request.url.params == httpx2.QueryParams(
logical_date=logical_date.isoformat(),
map_index="-1",
state="success",
)
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/task-instances/previous/test_dag/test_task":
assert request.url.params == httpx2.QueryParams(
map_index="0",
)
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/task-instances/previous/test_dag/test_task":
# Return None (null) when no previous TI found
return httpx2.Response(status_code=200, content="null")
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
nonlocal call_count
call_count += 1
if call_count < 2:
return httpx2.Response(status_code=500, json={"detail": "Internal Server Error"})
if request.url.path == "/variables/test_key":
return httpx2.Response(
status_code=200,
json={"key": "test_key", "value": "test_value"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/variables/non_existent_var":
return httpx2.Response(
status_code=404,
json={
"detail": {
"message": "Variable with key 'non_existent_var' not found",
"reason": "not_found",
}
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/variables/test_key":
return httpx2.Response(
status_code=500,
headers=[("content-Type", "application/json")],
json={
"reason": "internal_server_error",
"message": "Internal Server Error",
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=status_code, json={"detail": "Forbidden"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/variables/test_key":
return httpx2.Response(
status_code=201,
json={"message": "Variable successfully set"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.method == "DELETE" and request.url.path == "/variables/test_key":
return httpx2.Response(
status_code=200,
json={"count": 1},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
nonlocal call_count
call_count += 1
if call_count < 3:
return httpx2.Response(status_code=500, json={"detail": "Internal Server Error"})
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
return httpx2.Response(
status_code=201,
json={"key": "test_key", "value": value},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if (
request.url.path == "/xcoms/dag_id/run_id/task_id/key"
and request.url.params.get("map_index") == "2"
):
return httpx2.Response(
status_code=201,
json={"key": "test_key", "value": "test_value"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/xcoms/dag_id/run_id/task_id/key" and request.url.params.get(
"include_prior_dates"
):
return httpx2.Response(
status_code=201,
json={"key": "test_key", "value": "test_value"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
return httpx2.Response(
status_code=500,
headers=[("content-Type", "application/json")],
json={
"reason": "invalid_format",
"message": "XCom value is not a valid JSON",
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/xcoms/dag_id/run_id/task_id/key":
assert json.loads(request.read()) == values
return httpx2.Response(
status_code=201,
json={"message": "XCom successfully set"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=201,
json={"message": "XCom successfully set"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=201,
json={"message": "XCom successfully set"},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/connections/test_conn":
return httpx2.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 httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/connections/test_conn":
return httpx2.Response(
status_code=404,
json={
"reason": "not_found",
"message": "Connection with ID test_conn not found",
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
requests_seen.append(request)
return httpx2.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=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=status_code, json={"detail": "Forbidden"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=400, json={"detail": "Bad Request"})
return httpx2.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=httpx2.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: httpx2.Request) -> httpx2.Response:
params = request.url.params
assert params.get("partition_key") == "2024-01-01"
assert "partition_key_regexp_pattern" not in params
return httpx2.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(httpx2.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: httpx2.Request) -> httpx2.Response:
params = request.url.params
assert params.get("partition_key_regexp_pattern") == "^2024-01-"
assert "partition_key" not in params
return httpx2.Response(
status_code=200,
json={"asset_events": []},
)
client = make_client(httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=200,
json={"asset_events": []},
)
client = make_client(httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(
status_code=200,
json={"asset_events": []},
)
client = make_client(httpx2.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: httpx2.Request) -> httpx2.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 httpx2.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=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.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=httpx2.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: httpx2.Request) -> httpx2.Response:
params = request.url.params
assert "extra" not in params
return httpx2.Response(
status_code=200,
json={"asset_events": []},
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path in ("/assets/by-name", "/assets/by-uri"):
return httpx2.Response(
status_code=200,
json={
"name": "this_asset",
"uri": "s3://bucket/key",
"group": "asset",
"extra": {"foo": "bar"},
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path in ("/assets/by-name", "/assets/by-uri"):
return httpx2.Response(
status_code=404,
json={
"detail": {
"message": "Asset with name non_existent not found",
"reason": "not_found",
}
},
)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.url.path == "/assets/by-alias"
assert request.url.params["alias_name"] == "my_alias"
return httpx2.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=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(status_code=200, json=[])
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
return httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/test_trigger_conflict/test_run_id":
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/test_trigger_conflict_reset/test_run_id":
return httpx2.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 httpx2.Response(status_code=204)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/test_clear/test_run_id/clear":
return httpx2.Response(status_code=204)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/test_state/test_run_id/state":
return httpx2.Response(
status_code=200,
json={"state": "running"},
)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/count":
assert request.url.params["dag_id"] == "test_dag"
return httpx2.Response(status_code=200, json=1)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=200, json=2)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=200, json=2)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=200, json=2)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/previous":
assert request.url.params == httpx2.QueryParams(
dag_id="test_dag",
logical_date=logical_date.isoformat(),
)
# Return complete DagRun data
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/previous":
assert request.url.params == httpx2.QueryParams(
dag_id="test_dag",
logical_date=logical_date.isoformat(),
state="success",
)
# Return complete DagRun data
return httpx2.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 httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dag-runs/previous":
assert request.url.params == httpx2.QueryParams(
dag_id="test_dag",
logical_date=logical_date.isoformat(),
)
# Return None (null) when no previous Dag run found
return httpx2.Response(status_code=200, content="null")
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/task-reschedules/{ti_id}/start_date":
assert request.url.params.get("try_number") == "1"
return httpx2.Response(
status_code=200,
json="2024-01-01T00:00:00Z",
)
return httpx2.Response(status_code=422)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path in (f"/hitlDetails/{ti_id}"):
return httpx2.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 httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path in (f"/hitlDetails/{ti_id}"):
return httpx2.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 httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path in (f"/hitlDetails/{ti_id}"):
return httpx2.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 httpx2.Response(status_code=400, json={"detail": "Bad Request"})
client = make_client(transport=httpx2.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):
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):
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):
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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dags/test_dag":
return httpx2.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 httpx2.Response(status_code=200)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
requests_seen.append(request)
if request.url.path == "/variables/secret_key":
return httpx2.Response(
status_code=200,
json={"key": "secret_key", "value": "super-secret-value"},
)
return httpx2.Response(
status_code=404,
json={
"detail": {
"message": "The Dag was not found",
"reason": "not_found",
}
},
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dags/missing_dag":
return httpx2.Response(
status_code=404,
json={
"detail": {
"message": "The Dag with dag_id: `missing_dag` was not found",
"reason": "not_found",
}
},
)
return httpx2.Response(status_code=200)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == "/dags/test_dag":
return httpx2.Response(
status_code=500,
headers=[("content-Type", "application/json")],
json={
"reason": "internal_server_error",
"message": "Internal Server Error",
},
)
return httpx2.Response(status_code=200)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if request.url.path == f"/store/ti/{self.TI_ID}/job_id":
return httpx2.Response(status_code=200, json={"value": "spark_app_001"})
return httpx2.Response(status_code=400)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
requests_seen.append(request)
return httpx2.Response(status_code=200, json={"value": "spark_app_001"})
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(
status_code=404,
json={"detail": {"reason": "not_found", "message": "Task state key 'job_id' not found"}},
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
requests_seen.append(request)
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
body = json.loads(request.content)
assert body["value"] == "spark_app_001"
assert body["expires_at"] == "2026-05-21T12:00:00Z"
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
body = json.loads(request.content)
assert body.get("expires_at") is None
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.method == "DELETE"
assert request.url.path == f"/store/ti/{self.TI_ID}/job_id"
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
requests_seen.append(request)
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.method == "DELETE"
assert request.url.path == f"/store/ti/{self.TI_ID}"
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
if (
request.url.path == "/store/asset/by-name/value"
and request.url.params["name"] == "test_asset"
and request.url.params["key"] == "watermark"
):
return httpx2.Response(status_code=200, json={"value": "2026-04-30T00:00:00Z"})
return httpx2.Response(status_code=400)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=200, json={"value": "2026-04-30T00:00:00Z"})
return httpx2.Response(status_code=400)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(
status_code=404,
json={"detail": {"reason": "not_found", "message": "Asset state key 'watermark' not found"}},
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.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 httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.method == "DELETE"
assert request.url.path == "/store/asset/by-name/clear"
assert request.url.params["name"] == "test_asset"
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.method == "DELETE"
assert request.url.path == "/store/asset/by-uri/clear"
assert request.url.params["uri"] == "s3://bucket/key"
return httpx2.Response(status_code=204)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
assert request.method == "PATCH"
assert request.url.path == f"/callbacks/{callback_id}/run"
return httpx2.Response(
status_code=204,
headers={"Refreshed-API-Token": "execution-token"},
)
client = make_client(transport=httpx2.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: httpx2.Request) -> httpx2.Response:
return httpx2.Response(
status_code=409,
json={"detail": {"reason": "invalid_state", "previous_state": "running"}},
)
client = make_client(transport=httpx2.MockTransport(handle_request))
with pytest.raises(ServerResponseError) as err:
client.callbacks.run(callback_id)
assert err.value.response.status_code == 409