Repository: apache/airflow
File: providers/celery/tests/integration/celery/test_celery_executor.py (Python)

Code before:
                    raise AirflowException("fail")

        with _prepare_app(broker_url, execute=fake_execute) as app:
            executor = celery_executor.CeleryExecutor()
            assert executor.tasks == {}
            executor.start()

            with start_worker(app=app, logfile=sys.stdout, loglevel="info"):
                ti = workloads.TaskInstance.model_construct(
                    id=uuid7(),
                    task_id="success",
                    dag_id="id",
                    run_id="abc",
                    try_number=0,
                    priority_weight=1,
                    queue=celery_executor_utils.get_celery_configuration()["task_default_queue"],
                    executor_config=executor_config,
                )
                keys = [
                    TaskInstanceKey("id", "success", "abc", 0, -1),
                    TaskInstanceKey("id", "fail", "abc", 0, -1),
                ]
                for w in (
                    workloads.ExecuteTask.model_construct(
                        ti=ti,
                        token="token",
                        dag_rel_path="dag_rel_path",
                        bundle_info=None,
                        log_path="log_path",
                        dag_version_id=uuid7(),
                        pools_slot=123,
                    ),
                    workloads.ExecuteTask.model_construct(
                        ti=ti.model_copy(update={"task_id": "fail"}),
                        token="token",
                        dag_rel_path="dag_rel_path",
                        bundle_info=None,
                        log_path="log_path",
                        dag_version_id=uuid7(),
                        pools_slot=123,
                    ),
                ):
                    executor.queue_workload(w, session=None)

                executor.trigger_tasks(open_slots=10)
                for _ in range(20):
                    num_tasks = len(executor.tasks.keys())
                    if num_tasks == 2:
                        break
                    logger.info(
                        "Waiting 0.1 s for tasks to be processed asynchronously. Processed so far %d",

Review comment (suggestion):
This could work. Maybe we should use `make` rather than `model_construct` so that such as token generation can be from generated and using defaults

https://github.com/apache/airflow/blob/40f6ec1c6021f242e80e010043273d2a4cbd4887/airflow-core/src/airflow/executors/workloads.py#L111-L146


Code after:
            def fake_execute(input: str) -> None:
                """Fake execute_workload that parses JSON and fails for tasks with 'fail' in task_id."""
                import json

                workload_dict = json.loads(input)
                # Check if this is a task that should fail (task_id contains "fail")
                if "ti" in workload_dict and "task_id" in workload_dict["ti"]:
                    if "fail" in workload_dict["ti"]["task_id"]:
                        raise AirflowException("fail")
        else:
            # Airflow 2: execute_command receives command list
            def fake_execute(input: str) -> None:  # Use same parameter name as Airflow 3 version
                if "fail" in input:
                    raise AirflowException("fail")

        with _prepare_app(broker_url, execute=fake_execute) as app:
            executor = celery_executor.CeleryExecutor()
            assert executor.tasks == {}
            executor.start()

            with start_worker(app=app, logfile=sys.stdout, loglevel="info"):
                ti_success = workloads.TaskInstance.model_construct(
                    id=uuid7(),
                    task_id="success",
                    dag_id="id",
                    run_id="abc",
                    try_number=0,
                    priority_weight=1,
                    queue=celery_executor_utils.get_celery_configuration()["task_default_queue"],
                    executor_config=executor_config,
                )
                keys = [
                    TaskInstanceKey("id", "success", "abc", 0, -1),
                    TaskInstanceKey("id", "fail", "abc", 0, -1),
                ]
                dagrun = setup_dagrun_with_success_and_fail_tasks(dag_maker)
                ti_success, ti_fail = dagrun.task_instances
                for w in (
                    workloads.ExecuteTask.make(
                        ti=ti_success,
                    ),
                    workloads.ExecuteTask.make(ti=ti_fail),
                ):
                    executor.queue_workload(w, session=None)

                executor.trigger_tasks(open_slots=10)
                for _ in range(20):
                    num_tasks = len(executor.tasks.keys())
                    if num_tasks == 2:
                        break
                    logger.info(