This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ea349226ae0 Add Celery worker workload dispatch regression tests
(#65849)
ea349226ae0 is described below
commit ea349226ae01f1a530881a6eab46420ae946c91f
Author: Paul <[email protected]>
AuthorDate: Fri May 1 22:25:27 2026 -0700
Add Celery worker workload dispatch regression tests (#65849)
---
.../unit/celery/executors/test_celery_executor.py | 85 ++++++++++++++++++++++
1 file changed, 85 insertions(+)
diff --git
a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
index 23277c70e8f..309eba50cbd 100644
--- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
+++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
@@ -51,6 +51,7 @@ from tests_common.test_utils.version_compat import (
AIRFLOW_V_3_1_9_PLUS,
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_2_PLUS,
+ AIRFLOW_V_3_3_PLUS,
)
try:
@@ -884,6 +885,90 @@ def test_celery_tasks_registered_on_import():
)
[email protected](not AIRFLOW_V_3_0_PLUS, reason="execute_workload is only
used for Airflow 3+")
[email protected](AIRFLOW_V_3_3_PLUS, reason="pre-3.3 compatibility path
only applies before Airflow 3.3")
+def test_execute_workload_runs_execute_task_before_airflow_3_3():
+ """execute_workload routes serialized ExecuteTask payloads to supervise on
Airflow 3.0-3.2."""
+ from airflow.executors import workloads
+
+ workload = workloads.ExecuteTask(
+ ti=workloads.TaskInstance(
+ id="00000000-0000-0000-0000-000000000001",
+ dag_version_id="00000000-0000-0000-0000-000000000002",
+ task_id="test_task",
+ dag_id="test_dag",
+ run_id="test_run",
+ try_number=1,
+ map_index=-1,
+ pool_slots=1,
+ queue="default",
+ priority_weight=1,
+ ),
+ dag_rel_path="test_dag.py",
+ bundle_info=workloads.BundleInfo(name="test-bundle", version=None),
+ token="test-token",
+ log_path="test.log",
+ )
+ mock_current_task = mock.MagicMock()
+ mock_current_task.request.id = "test-celery-task-id"
+ mock_app = mock.MagicMock()
+ mock_app.current_task = mock_current_task
+
+ with (
+ mock.patch.object(celery_executor_utils, "app", mock_app),
+ mock.patch("airflow.sdk.execution_time.supervisor.supervise") as
mock_supervise,
+ ):
+
celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json())
+
+ mock_supervise.assert_called_once()
+ assert mock_supervise.call_args.kwargs["ti"].task_id == "test_task"
+ assert str(mock_supervise.call_args.kwargs["dag_rel_path"]) ==
"test_dag.py"
+ assert mock_supervise.call_args.kwargs["token"] == "test-token"
+ assert mock_supervise.call_args.kwargs["log_path"] == "test.log"
+
+
[email protected](not AIRFLOW_V_3_3_PLUS, reason="3.3+ path uses
BaseExecutor.run_workload")
+def test_execute_workload_runs_base_executor_workload_on_airflow_3_3_plus():
+ """execute_workload routes serialized ExecuteTask payloads to BaseExecutor
on Airflow 3.3+."""
+ from airflow.executors import workloads
+
+ workload = workloads.ExecuteTask(
+ ti=workloads.TaskInstance(
+ id="00000000-0000-0000-0000-000000000001",
+ dag_version_id="00000000-0000-0000-0000-000000000002",
+ task_id="test_task",
+ dag_id="test_dag",
+ run_id="test_run",
+ try_number=1,
+ map_index=-1,
+ pool_slots=1,
+ queue="default",
+ priority_weight=1,
+ ),
+ dag_rel_path="test_dag.py",
+ bundle_info=workloads.BundleInfo(name="test-bundle", version=None),
+ token="test-token",
+ log_path="test.log",
+ )
+ mock_current_task = mock.MagicMock()
+ mock_current_task.request.id = "test-celery-task-id"
+ mock_app = mock.MagicMock()
+ mock_app.current_task = mock_current_task
+
+ with (
+ mock.patch.object(celery_executor_utils, "app", mock_app),
+
mock.patch("airflow.executors.base_executor.BaseExecutor.run_workload") as
mock_run_workload,
+ ):
+
celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json())
+
+ mock_run_workload.assert_called_once()
+ decoded_workload = mock_run_workload.call_args.args[0]
+ assert decoded_workload.ti.task_id == "test_task"
+ assert str(decoded_workload.dag_rel_path) == "test_dag.py"
+ assert decoded_workload.token == "test-token"
+ assert decoded_workload.log_path == "test.log"
+
+
@pytest.mark.skipif(not AIRFLOW_V_3_1_9_PLUS, reason="TaskAlreadyRunningError
requires Airflow 3.1.9+")
@pytest.mark.skipif(
not hasattr(BaseExecutor, "run_workload"),
