Repository: apache/airflow
File: airflow-core/tests/unit/jobs/test_scheduler_job.py (Python)

Code before:
        b_tis_in_wrong_executor = [
            ti for ti in executor_to_tis.get(mock_executors[0], []) if ti.dag_id == "dag_b"
        ]
        assert len(b_tis_in_wrong_executor) == 0

    def task_helper(self, dag_maker, session, dag_id: str, task_num: int):
        dag_tasks = {}

        with dag_maker(dag_id=dag_id):
            for i in range(task_num):
                # Assign priority weight to certain tasks.
                if (i % 10) == 0:  # 10, 20, 30, 40, 50, ...
                    weight = int(i / 2)
                    dag_tasks[f"op{i}"] = EmptyOperator(task_id=f"dummy{i}", priority_weight=weight)
                else:
                    # No executor specified, runs on default executor
                    dag_tasks[f"op{i}"] = EmptyOperator(task_id=f"dummy{i}")

        dag_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)

        task_tis = {}

        tis_list = []
        for i in range(task_num):
            task_tis[f"ti{i}"] = dag_run.get_task_instance(dag_tasks[f"op{i}"].task_id, session)
            # add
            tis_list.append(task_tis[f"ti{i}"])

        for ti in tis_list:
            ti.state = State.SCHEDULED
            ti.dag_model.max_active_tasks = 4

        session.flush()

        return tis_list

    @conf_vars(
        {
            ("scheduler", "max_tis_per_query"): "100",
            ("scheduler", "max_dagruns_to_create_per_loop"): "10",
            ("scheduler", "max_dagruns_per_loop_to_schedule"): "20",
            ("core", "parallelism"): "100",
            ("core", "max_active_tasks_per_dag"): "4",
            ("core", "max_active_runs_per_dag"): "10",
            ("core", "default_pool_task_slot_count"): "64",
        }
    )
    def test_per_dr_limit_applied_in_task_query(self, dag_maker, mock_executors):
        scheduler_job = Job()
        scheduler_job.executor.parallelism = 100
        scheduler_job.executor.slots_available = 70

Review comment (suggestion):
[nitpick] The comment "# add" is too vague and doesn't add meaningful information. Consider either removing it or making it more descriptive, such as "# Add task instance to list for state updates".
```suggestion
            # Add task instance to list for state updates
```

Code after:
        ti4 = dag_run.get_task_instance(op4.task_id, session)
        ti5 = dag_run.get_task_instance(op5.task_id, session)

        tis_tuple = (ti1, ti2, ti3, ti4, ti5)
        for ti in tis_tuple:
            ti.state = State.SCHEDULED

        session.flush()

        res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)

        assert len(res) == 5
        res_ti_keys = [res_ti.key for res_ti in res]
        for ti in tis_tuple:
            assert ti.key in res_ti_keys

    @conf_vars({("core", "multi_team"): "true"})
    def test_find_executable_task_instances_executor_with_teams(self, dag_maker, mock_executors, session):
        """
        Test that tasks are correctly routed to team-specific executors when multi-team is enabled
        """
        clear_db_teams()
        clear_db_dag_bundles()

        team1 = Team(name="team_a")
        team2 = Team(name="team_b")
        session.add_all([team1, team2])
        session.flush()

        bundle1 = DagBundleModel(name="bundle_a")
        bundle2 = DagBundleModel(name="bundle_b")
        bundle1.teams.append(team1)
        bundle2.teams.append(team2)
        session.add_all([bundle1, bundle2])
        session.flush()

        mock_executors[0].team_name = "team_a"
        mock_executors[1].team_name = "team_b"

        with dag_maker(dag_id="dag_a", bundle_name="bundle_a", session=session):
            op1 = EmptyOperator(task_id="task_a_default")  # No explicit executor - should use team's default
            op2 = EmptyOperator(
                task_id="task_a_explicit", executor="default_exec"
            )  # Team-specific explicit executor
        dr1 = dag_maker.create_dagrun()

        with dag_maker(dag_id="dag_b", bundle_name="bundle_b", session=session):
            op3 = EmptyOperator(task_id="task_b_default")  # Team b's default
            op4 = EmptyOperator(task_id="task_b_explicit", executor="secondary_exec")  # Team b explicit
        dr2 = dag_maker.create_dagrun()