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

Code before:
            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
        scheduler_job.max_tis_per_query = 100
        self.job_runner = SchedulerJobRunner(job=scheduler_job)
        session = settings.Session()

        # Use the same run_id.
        task_maker(dag_maker, session, "dag_1300_tasks", 1300, 4, "run1")
        task_maker(dag_maker, session, "dag_1200_tasks", 1200, 4, "run1")
        task_maker(dag_maker, session, "dag_1100_tasks", 1100, 4, "run1")
        task_maker(dag_maker, session, "dag_100_tasks", 100, 4, "run1")

Review comment (refactor):
The `task_helper` method (lines 1229-1258) is nearly identical to the module-level `task_maker` function (lines 184-227). The only significant differences are:
1. `task_helper` has a hardcoded `max_active_tasks = 4` (line 1254) instead of accepting it as a parameter
2. `task_helper` doesn't support custom `run_id` parameter

Consider removing this duplicate method and using the `task_maker` function instead, or refactor to reduce code duplication. For example, the test at line 1271 could use `task_maker` directly.

Code after:
        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()

        # DAG with no team (global)
        with dag_maker(dag_id="dag_global", session=session):
            op5 = EmptyOperator(task_id="task_global")  # Global task - any executor
        dr3 = dag_maker.create_dagrun()

        tis = [
            dr1.get_task_instance(op1.task_id, session),
            dr1.get_task_instance(op2.task_id, session),
            dr2.get_task_instance(op3.task_id, session),