ashb commented on code in PR #44723:
URL: https://github.com/apache/airflow/pull/44723#discussion_r1873348994
##########
task_sdk/tests/api/test_client.py:
##########
@@ -82,6 +85,88 @@ def make_client(transport: httpx.MockTransport) -> Client:
return Client(base_url="test://server", token="", transport=transport)
+class TestTaskInstanceOperations:
+ """
+ Test that the TestVariableOperations 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):
+ # Simulate a successful response from the server that starts a task
+ ti_id = uuid6.uuid7()
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/state":
+ return httpx.Response(
+ status_code=204,
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad
Request"})
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ try:
+ client.task_instances.start(ti_id, 100, "2024-10-31T12:00:00Z")
+ except Exception as e:
+ pytest.fail(f"Unexpected error occurred: {e}")
+
+ @pytest.mark.parametrize("state", [state for state in TerminalTIState])
+ def test_task_instance_finish(self, state):
+ # Simulate a successful response from the server that finishes a task
+ ti_id = uuid6.uuid7()
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/state":
+ return httpx.Response(
+ status_code=204,
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad
Request"})
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ try:
+ client.task_instances.finish(ti_id, state=state,
when="2024-10-31T12:00:00Z")
+ except Exception as e:
+ pytest.fail(f"Unexpected error occurred: {e}")
+
+ def test_task_instance_heartbeat(self):
+ # Simulate a successful response from the server is used to heartbeat
+ ti_id = uuid6.uuid7()
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/heartbeat":
+ return httpx.Response(
+ status_code=204,
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad
Request"})
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ try:
+ client.task_instances.heartbeat(ti_id, 100)
+ except Exception as e:
+ pytest.fail(f"Unexpected error occurred: {e}")
+
+ def test_task_instance_defer(self):
+ # Simulate a successful response from the server that defers a task
+ ti_id = uuid6.uuid7()
+
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/task-instances/{ti_id}/state":
+ return httpx.Response(
+ status_code=204,
+ )
+ return httpx.Response(status_code=400, json={"detail": "Bad
Request"})
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ try:
+ msg = DeferTask(
+
classpath="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
+ trigger_kwargs={"moment": "2024-11-07T12:34:59Z",
"end_from_trigger": False},
+ next_method="execute_complete",
+ )
+ client.task_instances.defer(ti_id, msg)
+ except Exception as e:
+ pytest.fail(f"Unexpected error occurred: {e}")
Review Comment:
This isn't needed, pytest will already mark it as failed if the test fails
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
