amoghrajesh commented on code in PR #66160:
URL: https://github.com/apache/airflow/pull/66160#discussion_r3186762333
##########
task-sdk/tests/task_sdk/api/test_client.py:
##########
@@ -1700,3 +1702,188 @@ def handle_request(request: httpx.Request) ->
httpx.Response:
with pytest.raises(ServerResponseError):
client.dags.get(dag_id="test_dag")
+
+
+class TestTaskStateOperations:
+ TI_ID = uuid7()
+
+ def test_get_success(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ if request.url.path == f"/state/ti/{self.TI_ID}/job_id":
+ return httpx.Response(status_code=200, json={"value":
"spark_app_001"})
+ return httpx.Response(status_code=400)
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ result = client.task_state.get(ti_id=self.TI_ID, key="job_id")
+
+ assert isinstance(result, TaskStateResponse)
+ assert result.value == "spark_app_001"
+
+ def test_get_raises_on_404(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ return httpx.Response(
+ status_code=404,
+ json={"detail": {"reason": "not_found", "message": "Task state
key 'job_id' not found"}},
+ )
+
+ client = make_client(transport=httpx.MockTransport(handle_request))
+ with pytest.raises(ServerResponseError):
+ client.task_state.get(ti_id=self.TI_ID, key="job_id")
Review Comment:
Fixed it
--
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]
