Fix provider executor tests broken in main#67268
Conversation
anishgirianish
requested review from
hussein-awala,
jedcunningham,
jscheffl and
o-nikolas
as code owners
boring-cyborg
Bot
added
area:providers
provider:amazon
anishgirianish
force-pushed
the
fix-provider-executor-tests-key-types
branch
from
0fb69dd to
b8d0897
Compare
amoghrajesh
left a comment
amoghrajesh
left a comment
There was a problem hiding this comment.
Review findings
F1 — BLOCKING: Top-level airflow.models.callback import breaks compat tests
The PR adds from airflow.models.callback import CallbackKey as a top-level import in test_lambda_executor.py. The airflow.models.callback module was introduced in Airflow 3.3 and doesn't exist in 2.11.1, 3.0.6, or 3.1.8. The compat CI jobs fail at collection time with:
ModuleNotFoundError: No module named 'airflow.models.callback'
The test that uses CallbackKey (test_task_sdk_callback) is correctly guarded with @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, ...), but the top-level import fires unconditionally and kills the whole test module on older versions.
Fix: Guard the import:
if AIRFLOW_V_3_3_PLUS:
from airflow.models.callback import CallbackKeyF2 — BLOCKING: CallbackKey(id="callback_123") is invalid Python
CallbackKey in the current codebase is CallbackKey = str (a plain type alias, airflow/models/callback.py:41). Calling str(id="callback_123") raises TypeError — str() doesn't accept id= as a keyword argument.
Since AIRFLOW_V_3_3_PLUS is now True on main (version 3.3.0), this test is no longer skipped and will fail when run.
Fix: Use callback_id = "callback_123" (the old form was correct), or CallbackKey("callback_123") if you want the type to be explicit.
F3 — NOTE: MyPy and Celery low-dep failures appear pre-existing
The mypy failure is in providers/google (CreateCachedContentConfig contents type mismatch) and the Celery test_process_workloads_routes_execute_callback failure are in untouched code — worth confirming both fail on main before merging so they don't accidentally land on your shoulders.
What's good
- Replacing
Mock(spec=tuple)/Mock(spec=list)/ bare strings withMock(spec=TaskInstanceKey)or realTaskInstanceKey(...)instances is the right fix — these now satisfy thestate_class_for_keyguard. ser_workload_key = str(workload_key)in the Lambda executor is the correct defensive serialization.- The K8s test fix (
TaskInstanceKey("dag", "task", "run_id", 1, -1)replacing a raw tuple) is correct and overdue. - Narrowing
execute_asynctoBatchJobWorkloadKeyis accurate.
F1 and F2 are both in the same test function, so the fix should be a small follow-up commit.
Drafted-by: Claude Code (claude-sonnet-4-6); reviewed by @amoghrajesh before posting
anishgirianish
force-pushed
the
fix-provider-executor-tests-key-types
branch
from
a11d16a to
56ad11d
Compare
amoghrajesh
left a comment
amoghrajesh
left a comment
There was a problem hiding this comment.
Yeah I think its good now, kets see what the CI says.
Switch from Mock(spec=tuple) / bare-string callback key to Mock(spec=TaskInstanceKey) and a real CallbackKey(...) instance, matching the pattern established in apache#67268 for the rest of the ECS/Lambda/Batch executor tests. CallbackKey became a frozen dataclass in apache#66973 and no longer accepts bare strings; this test was added in this PR so it was missed by the apache#67268 cleanup sweep.
| executor._process_workloads([workload]) | ||
|
|
||
| mock_send_workloads.assert_called_once_with([(callback_id, workload, expected_queue, None)]) | ||
| mock_send_workloads.assert_called_once_with([(workload.callback.key, workload, expected_queue, None)]) |
There was a problem hiding this comment.
FWIW, this could just be workload.key since the ExecuteCallback.key property returns self.callback.key. But thanks for fixing this. I'm curious why my PR was green but caused all these issues. Sorry about that.
* Add ExecuteCallback support to AWS ECS Executor Enables the ECS executor to dispatch ExecuteCallback workloads (deadline alerts) alongside regular ExecuteTask workloads. Builds on #65392 which widened BaseExecutor signatures to accept WorkloadKey. - supports_callbacks = True (gated on AIRFLOW_V_3_3_PLUS) - Widen key types to WorkloadKey throughout EcsQueuedTask / EcsTaskCollection - Branch _process_workloads on ExecuteTask vs ExecuteCallback - Add AIRFLOW_V_3_3_PLUS to version_compat.py - Unit tests for queueing, processing, serialization, sync, mixed keys * Rename task-named methods/attrs and fix older-Airflow compat Renames (mirrors the merged Lambda callback PR — straight rename, no shim, executor-internal surface): sync_running_tasks -> sync_running_workloads attempt_task_runs -> attempt_workload_runs pending_tasks (attr) -> pending_workloads __update_running_task -> __update_running_workload __handle_failed_task -> __handle_failed_workload Fix CI on older Airflow compat tests: - Restore queue_workload() override. Airflow 3.3+ BaseExecutor routes ExecuteCallback natively, but pre-3.3 raises ValueError for anything not ExecuteTask. Override works across versions. - Import AIRFLOW_V_3_3_PLUS from tests_common (main bumped to 3.3). check-airflow-v-imports-in-tests hook disallows provider-internal version_compat imports from test files. * Apply suggestions from code review Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com> * Widen CommandType and pre-declare loop locals to drop type ignores Mirror the Lambda executor pattern: widen CommandType to Sequence[str] | Sequence[ExecuteTask | ExecuteCallback] on Airflow 3.3+, and pre-declare queue/key/command at the top of the _process_workloads loop so mypy doesn't infer narrow types from the first if-branch. Removes five # type: ignore comments that previously covered the signature mismatch and cross-branch reassignment. * Fix docs autoapi duplicate-object warnings Two adjustments to silence Sphinx duplicate-object warnings between the Lambda and ECS executors that block the docs build: - Add `:sphinx-autoapi-skip:` docstring on ECS `_process_workloads` so autoapi skips this private method. - Split the inner union in ECS `CommandType` so its rendered TypeAlias text differs from Lambda's, avoiding a duplicate. The split form is also stricter (disallows mixed task+callback lists, which our code never produces). * Add trailing period to sphinx-autoapi-skip docstring for D415 * Use TaskInstanceKey and CallbackKey in test_collection_mixed_key_types Switch from Mock(spec=tuple) / bare-string callback key to Mock(spec=TaskInstanceKey) and a real CallbackKey(...) instance, matching the pattern established in #67268 for the rest of the ECS/Lambda/Batch executor tests. CallbackKey became a frozen dataclass in #66973 and no longer accepts bare strings; this test was added in this PR so it was missed by the #67268 cleanup sweep. --------- Co-authored-by: D. Ferruzzi <ferruzzi@amazon.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
3 participants
Was generative AI tooling used to co-author this PR?
Several amazon (batch/ecs/lambda) and cncf/kubernetes executor tests pass
Mock(spec=tuple), tuple literals, or bare strings as workload keys, which now fail withTypeError: Unknown workload key typeafter thestate_class_for_keyguard was tightened.Swaps them for
Mock(spec=TaskInstanceKey)/ realTaskInstanceKey(...)/CallbackKey(...). Also fixes the Lambda executor's callback serialization which fed aCallbackKeydataclass intojson.dumps, now usesstr(workload_key).related: #66973
cc: @ferruzzi
{pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.