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

Code before:
                    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.make(
                        ti=ti,
                        dag_rel_path="dag_rel_path",
                        generator=None,
                        bundle_info=None,
                        sentry_integration="",
                    ),
                    workloads.ExecuteTask.make(
                        ti=ti.model_copy(update={"task_id": "fail"}),
                        dag_rel_path="dag_rel_path",
                        generator=None,
                        bundle_info=None,
                        sentry_integration="",
                    ),
                ):
                    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",
                        num_tasks,
                    )
                    sleep(0.4)
                assert list(executor.tasks.keys()) == keys
                assert executor.event_buffer[keys[0]][0] == State.QUEUED
                assert executor.event_buffer[keys[1]][0] == State.QUEUED

                executor.end(synchronous=True)

        assert executor.event_buffer[keys[0]][0] == State.SUCCESS
        assert executor.event_buffer[keys[1]][0] == State.FAILED

        assert keys[0] not in executor.tasks
        assert keys[1] not in executor.tasks

        assert executor.queued_tasks == {}

    def test_error_sending_task(self):
        from airflow.providers.celery.executors import celery_executor, celery_executor_utils

Review comment (suggestion):
Yes seems like the right thing to do, actual executors do `make` cos `model_construct` bypasses validation: https://docs.pydantic.dev/latest/api/base_model/#pydantic.BaseModel.model_construct



Code after:
        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(
                        "Waiting 0.1 s for tasks to be processed asynchronously. Processed so far %d",
                        num_tasks,
                    )
                    sleep(0.4)
                assert list(executor.tasks.keys()) == keys
                assert executor.event_buffer[keys[0]][0] == State.QUEUED
                assert executor.event_buffer[keys[1]][0] == State.QUEUED

                executor.end(synchronous=True)

        assert executor.event_buffer[keys[0]][0] == State.SUCCESS
        assert executor.event_buffer[keys[1]][0] == State.FAILED

        assert keys[0] not in executor.tasks
        assert keys[1] not in executor.tasks