Fix task sdk client dry-run mode (#46524)

* Fix task sdk client dry-run mode The `run_after` field was missing in the fake dag run response. * Fix trailing whitespaces
apache · Feb 6, 2025 · 7af4571 · 7af4571
1 parent 993a009
commit 7af4571
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/task_sdk/src/airflow/sdk/api/client.py b/task_sdk/src/airflow/sdk/api/client.py
@@ -344,6 +344,7 @@ def noop_handler(request: httpx.Request) -> httpx.Response:
                     "logical_date": "2021-01-01T00:00:00Z",
                     "start_date": "2021-01-01T00:00:00Z",
                     "run_type": DagRunType.MANUAL,
+                    "run_after": "2021-01-01T00:00:00Z",
                 },
                 "max_tries": 0,
             },

diff --git a/task_sdk/tests/api/test_client.py b/task_sdk/tests/api/test_client.py
@@ -37,6 +37,11 @@ def make_client(transport: httpx.MockTransport) -> Client:
     return Client(base_url="test://server", token="", transport=transport)
 
 
+def make_client_w_dry_run() -> Client:
+    """Get a client with dry_run enabled"""
+    return Client(base_url=None, dry_run=True, token="")
+
+
 def make_client_w_responses(responses: list[httpx.Response]) -> Client:
     """Helper fixture to create a mock client with custom responses."""
 
@@ -49,6 +54,34 @@ def handle_request(request: httpx.Request) -> httpx.Response:
 
 
 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",
+                    },
+                    "max_tries": 0,
+                },
+            ),
+        ],
+    )
+    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
+
     def test_error_parsing(self):
         responses = [
             httpx.Response(422, json={"detail": [{"loc": ["#0"], "msg": "err", "type": "required"}]})