Skip to content

Fix TaskSDK Tests #46496

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bugraoz93 wants to merge 2 commits into from
Closed

Fix TaskSDK Tests #46496

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3f9adb1
Fix task_sdk tests
bugraoz93 Feb 5, 2025
988688b
Fix task_sdk tests and mypy
bugraoz93 Feb 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions task_sdk/src/airflow/sdk/api/datamodels/_generated.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ class DagRun(BaseModel):
logical_date: Annotated[datetime, Field(title="Logical Date")]
data_interval_start: Annotated[datetime | None, Field(title="Data Interval Start")] = None
data_interval_end: Annotated[datetime | None, Field(title="Data Interval End")] = None
run_after: Annotated[datetime, Field(title="Run After")]
start_date: Annotated[datetime, Field(title="Start Date")]
end_date: Annotated[datetime | None, Field(title="End Date")] = None
run_type: DagRunType
Expand Down
12 changes: 11 additions & 1 deletion task_sdk/tests/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,10 @@ def __call__(
logical_date: str | datetime = ...,
data_interval_start: str | datetime = ...,
data_interval_end: str | datetime = ...,
run_after: str | datetime = ...,
start_date: str | datetime = ...,
run_type: str = ...,
conf=...,
) -> TIRunContext: ...


Expand All @@ -163,11 +165,13 @@ def __call__(
self,
dag_id: str = ...,
run_id: str = ...,
logical_date: str = ...,
logical_date: str | datetime = ...,
data_interval_start: str | datetime = ...,
data_interval_end: str | datetime = ...,
run_after: str | datetime = ...,
start_date: str | datetime = ...,
run_type: str = ...,
conf=...,
) -> dict[str, Any]: ...


Expand All @@ -182,6 +186,7 @@ def _make_context(
logical_date: str | datetime = "2024-12-01T01:00:00Z",
data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
run_after: str | datetime = "2024-12-01T01:00:00Z",
start_date: str | datetime = "2024-12-01T01:00:00Z",
run_type: str = "manual",
conf=None,
Expand All @@ -193,6 +198,7 @@ def _make_context(
logical_date=logical_date, # type: ignore
data_interval_start=data_interval_start, # type: ignore
data_interval_end=data_interval_end, # type: ignore
run_after=run_after, # type: ignore
start_date=start_date, # type: ignore
run_type=run_type, # type: ignore
conf=conf,
Expand All @@ -213,17 +219,21 @@ def _make_context_dict(
logical_date: str | datetime = "2024-12-01T00:00:00Z",
data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
run_after: str | datetime = "2024-12-01T00:00:00Z",
start_date: str | datetime = "2024-12-01T00:00:00Z",
run_type: str = "manual",
conf=None,
) -> dict[str, Any]:
context = make_ti_context(
dag_id=dag_id,
run_id=run_id,
logical_date=logical_date,
data_interval_start=data_interval_start,
data_interval_end=data_interval_end,
run_after=run_after,
start_date=start_date,
run_type=run_type,
conf=conf,
)
return context.model_dump(exclude_unset=True, mode="json")

Expand Down
2 changes: 2 additions & 0 deletions task_sdk/tests/execution_time/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ def _create_task_instance(
logical_date: str | datetime = "2024-12-01T01:00:00Z",
data_interval_start: str | datetime = "2024-12-01T00:00:00Z",
data_interval_end: str | datetime = "2024-12-01T01:00:00Z",
run_after: str | datetime = "2024-12-01T01:00:00Z",
start_date: str | datetime = "2024-12-01T01:00:00Z",
run_type: str = "manual",
try_number: int = 1,
Expand All @@ -146,6 +147,7 @@ def _create_task_instance(
logical_date=logical_date,
data_interval_start=data_interval_start,
data_interval_end=data_interval_end,
run_after=run_after,
start_date=start_date,
run_type=run_type,
conf=conf,
Expand Down
11 changes: 5 additions & 6 deletions task_sdk/tests/execution_time/test_supervisor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ def subprocess_main():
dag_rel_path=os.devnull,
bundle_info=FAKE_BUNDLE,
what=TaskInstance(
id="4d828a62-a417-4936-a7a6-2b3fabacecab",
id=TI_ID,
task_id="b",
dag_id="c",
run_id="d",
Expand Down Expand Up @@ -196,7 +196,7 @@ def subprocess_main():
dag_rel_path=os.devnull,
bundle_info=FAKE_BUNDLE,
what=TaskInstance(
id="4d828a62-a417-4936-a7a6-2b3fabacecab",
id=TI_ID,
task_id="b",
dag_id="c",
run_id="d",
Expand All @@ -220,7 +220,7 @@ def subprocess_main():
dag_rel_path=os.devnull,
bundle_info=FAKE_BUNDLE,
what=TaskInstance(
id=uuid7(),
id=TI_ID,
task_id="b",
dag_id="c",
run_id="d",
Expand Down Expand Up @@ -251,13 +251,12 @@ def subprocess_main():
print("output", flush=True)
sleep(0.05)

ti_id = uuid7()
spy = spy_agency.spy_on(sdk_client.TaskInstanceOperations.heartbeat)
proc = ActivitySubprocess.start(
dag_rel_path=os.devnull,
bundle_info=FAKE_BUNDLE,
what=TaskInstance(
id=ti_id,
id=TI_ID,
task_id="b",
dag_id="c",
run_id="d",
Expand All @@ -267,7 +266,7 @@ def subprocess_main():
target=subprocess_main,
)
assert proc.wait() == 0
assert spy.called_with(ti_id, pid=proc.pid) # noqa: PGH005
assert spy.called_with(TI_ID, pid=proc.pid) # noqa: PGH005
# The exact number we get will depend on timing behaviour, so be a little lenient
assert 1 <= len(spy.calls) <= 4

Expand Down
3 changes: 2 additions & 1 deletion task_sdk/tests/execution_time/test_task_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,8 @@ def test_recv_StartupDetails(self):
b'"id": "4d828a62-a417-4936-a7a6-2b3fabacecab", "task_id": "a", "try_number": 1, "run_id": "b", "dag_id": "c" }, '
b'"ti_context":{"dag_run":{"dag_id":"c","run_id":"b","logical_date":"2024-12-01T01:00:00Z",'
b'"data_interval_start":"2024-12-01T00:00:00Z","data_interval_end":"2024-12-01T01:00:00Z",'
b'"start_date":"2024-12-01T01:00:00Z","end_date":null,"run_type":"manual","conf":null},'
b'"run_after":"2024-12-01T00:00:00Z", "start_date":"2024-12-01T01:00:00Z","end_date":null,'
b'"run_type":"manual","conf":null},'
b'"max_tries":0,"variables":null,"connections":null},"file": "/dev/null", "dag_rel_path": "/dev/null", "bundle_info": {"name": '
b'"any-name", "version": "any-version"}, "requests_fd": '
+ str(w2.fileno()).encode("ascii")
Expand Down
Loading