You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/04/11 14:30:24 UTC
(airflow) branch main updated: Improvre `TestHandleRequest` to catch more bugs (#49125)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ad452089774 Improvre `TestHandleRequest` to catch more bugs (#49125)
ad452089774 is described below
commit ad4520897748d2ad17bcb468d56754e68b6af4ab
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Fri Apr 11 20:00:18 2025 +0530
Improvre `TestHandleRequest` to catch more bugs (#49125)
We would have caught the bug fixed in https://github.com/apache/airflow/pull/49095 (non triggerer one), with this change.
---
task-sdk/src/airflow/sdk/api/client.py | 20 ++++++++++----------
task-sdk/tests/task_sdk/api/test_client.py | 10 +++++-----
.../tests/task_sdk/execution_time/test_supervisor.py | 15 +++++++++------
3 files changed, 24 insertions(+), 21 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
index 8f2e8dba05a..df7679cab87 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -186,13 +186,13 @@ class TaskInstanceOperations:
body = TISkippedDownstreamTasksStatePayload(tasks=msg.tasks)
self.client.patch(f"task-instances/{id}/skip-downstream", content=body.model_dump_json())
- def set_rtif(self, id: uuid.UUID, body: dict[str, str]) -> dict[str, bool]:
+ def set_rtif(self, id: uuid.UUID, body: dict[str, str]) -> OKResponse:
"""Set Rendered Task Instance Fields via the API server."""
self.client.put(f"task-instances/{id}/rtif", json=body)
# Any error from the server will anyway be propagated down to the supervisor,
# so we choose to send a generic response to the supervisor over the server response to
# decouple from the server response string
- return {"ok": True}
+ return OKResponse(ok=True)
def get_previous_successful_dagrun(self, id: uuid.UUID) -> PrevSuccessfulDagRunResponse:
"""
@@ -279,14 +279,14 @@ class VariableOperations:
raise
return VariableResponse.model_validate_json(resp.read())
- def set(self, key: str, value: str | None, description: str | None = None):
+ def set(self, key: str, value: str | None, description: str | None = None) -> OKResponse:
"""Set an Airflow Variable via the API server."""
body = VariablePostBody(val=value, description=description)
self.client.put(f"variables/{key}", content=body.model_dump_json())
# Any error from the server will anyway be propagated down to the supervisor,
# so we choose to send a generic response to the supervisor over the server response to
# decouple from the server response string
- return {"ok": True}
+ return OKResponse(ok=True)
class XComOperations:
@@ -353,7 +353,7 @@ class XComOperations:
value,
map_index: int | None = None,
mapped_length: int | None = None,
- ) -> dict[str, bool]:
+ ) -> OKResponse:
"""Set a XCom value via the API server."""
# TODO: check if we need to use map_index as params in the uri
# ref: https://github.com/apache/airflow/blob/v2-10-stable/airflow/api_connexion/openapi/v1.yaml#L1785C1-L1785C81
@@ -366,7 +366,7 @@ class XComOperations:
# Any error from the server will anyway be propagated down to the supervisor,
# so we choose to send a generic response to the supervisor over the server response to
# decouple from the server response string
- return {"ok": True}
+ return OKResponse(ok=True)
def delete(
self,
@@ -375,7 +375,7 @@ class XComOperations:
task_id: str,
key: str,
map_index: int | None = None,
- ) -> dict[str, bool]:
+ ) -> OKResponse:
"""Delete a XCom with given key via the API server."""
params = {}
if map_index is not None and map_index >= 0:
@@ -384,7 +384,7 @@ class XComOperations:
# Any error from the server will anyway be propagated down to the supervisor,
# so we choose to send a generic response to the supervisor over the server response to
# decouple from the server response string
- return {"ok": True}
+ return OKResponse(ok=True)
class AssetOperations:
@@ -453,7 +453,7 @@ class DagRunOperations:
conf: dict | None = None,
logical_date: datetime | None = None,
reset_dag_run: bool = False,
- ):
+ ) -> OKResponse | ErrorResponse:
"""Trigger a DAG run via the API server."""
body = TriggerDAGRunPayload(logical_date=logical_date, conf=conf or {}, reset_dag_run=reset_dag_run)
@@ -474,7 +474,7 @@ class DagRunOperations:
return OKResponse(ok=True)
- def clear(self, dag_id: str, run_id: str):
+ def clear(self, dag_id: str, run_id: str) -> OKResponse:
"""Clear a DAG run via the API server."""
self.client.post(f"dag-runs/{dag_id}/{run_id}/clear")
# TODO: Error handling
diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py
index 8e4ad0cab71..c6267496029 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -420,7 +420,7 @@ class TestTaskInstanceOperations:
client = make_client(transport=httpx.MockTransport(handle_request))
result = client.task_instances.set_rtif(id=TI_ID, body=rendered_fields)
- assert result == {"ok": True}
+ assert result == OKResponse(ok=True)
def test_get_count_basic(self):
"""Test basic get_count functionality with just dag_id."""
@@ -555,7 +555,7 @@ class TestVariableOperations:
client = make_client(transport=httpx.MockTransport(handle_request))
result = client.variables.set(key="test_key", value="test_value", description="test_description")
- assert result == {"ok": True}
+ assert result == OKResponse(ok=True)
class TestXCOMOperations:
@@ -706,7 +706,7 @@ class TestXCOMOperations:
key="key",
value=values,
)
- assert result == {"ok": True}
+ 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
@@ -731,7 +731,7 @@ class TestXCOMOperations:
value="value1",
map_index=2,
)
- assert result == {"ok": True}
+ 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
@@ -758,7 +758,7 @@ class TestXCOMOperations:
map_index=2,
mapped_length=3,
)
- assert result == {"ok": True}
+ assert result == OKResponse(ok=True)
class TestConnectionOperations:
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 77d052fe017..f8baf2e03ee 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -1053,7 +1053,7 @@ class TestHandleRequest:
"variables.set",
("test_key", "test_value", "test_description"),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="set_variable",
),
pytest.param(
@@ -1147,7 +1147,7 @@ class TestHandleRequest:
None,
),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="set_xcom",
),
pytest.param(
@@ -1171,7 +1171,7 @@ class TestHandleRequest:
None,
),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="set_xcom_with_map_index",
),
pytest.param(
@@ -1196,7 +1196,7 @@ class TestHandleRequest:
3,
),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="set_xcom_with_map_index_and_mapped_length",
),
pytest.param(
@@ -1217,7 +1217,7 @@ class TestHandleRequest:
2,
),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="delete_xcom",
),
# we aren't adding all states under TerminalTIState here, because this test's scope is only to check
@@ -1246,7 +1246,7 @@ class TestHandleRequest:
"task_instances.set_rtif",
(TI_ID, {"field1": "rendered_value1", "field2": "rendered_value2"}),
{},
- {"ok": True},
+ OKResponse(ok=True),
id="set_rtif",
),
pytest.param(
@@ -1525,6 +1525,9 @@ class TestHandleRequest:
# and deserialize it to the correct message type
# Only decode the buffer if it contains data. An empty buffer implies no response was written.
+ if not val and (mock_response and not isinstance(mock_response, OKResponse)):
+ pytest.fail("Expected a response, but got an empty buffer.")
+
if val:
# Using BytesIO to simulate a readable stream for CommsDecoder.
input_stream = BytesIO(val)