diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 1d1db8cfb0739..96174734062f1 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -283,7 +283,7 @@ def get_previous_successful_dagrun(self, id: uuid.UUID) -> PrevSuccessfulDagRunR def get_reschedule_start_date(self, id: uuid.UUID, try_number: int = 1) -> TaskRescheduleStartDate: """Get the start date of a task reschedule via the API server.""" resp = self.client.get(f"task-reschedules/{id}/start_date", params={"try_number": try_number}) - return TaskRescheduleStartDate.model_construct(start_date=resp.json()) + return TaskRescheduleStartDate(start_date=resp.json()) def get_count( self, diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index fd52433488ac1..75f6254974722 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1249,7 +1249,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: result = client.task_instances.get_reschedule_start_date(id=ti_id, try_number=1) assert isinstance(result, TaskRescheduleStartDate) - assert result.start_date == "2024-01-01T00:00:00Z" + assert result.start_date == datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc) class TestHITLOperations: