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

Code before:
        scheduler_job = Job()
        self.job_runner = SchedulerJobRunner(job=scheduler_job)
        session = settings.Session()

        dr = dag_maker.create_dagrun()

        ti = dr.task_instances[0]
        ti.state = State.SCHEDULED
        session.merge(ti)
        session.commit()

        res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)
        session.flush()
        assert len(res) == 0
        session.rollback()

    def test_infinite_pool(self, dag_maker):
        dag_id = "SchedulerJobTest.test_infinite_pool"
        with dag_maker(dag_id=dag_id, max_active_tasks=16):
            EmptyOperator(task_id="dummy", pool="infinite_pool")

        scheduler_job = Job()
        self.job_runner = SchedulerJobRunner(job=scheduler_job)
        session = settings.Session()

        dr = dag_maker.create_dagrun()
        ti = dr.task_instances[0]
        ti.state = State.SCHEDULED
        session.merge(ti)
        infinite_pool = Pool(
            pool="infinite_pool",
            slots=-1,
            description="infinite pool",
            include_deferred=False,
        )
        session.add(infinite_pool)
        session.commit()

        res = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session)
        session.flush()
        assert len(res) == 1
        session.rollback()

    def test_not_enough_pool_slots(self, caplog, dag_maker):
        dag_id = "SchedulerJobTest.test_test_not_enough_pool_slots"
        with dag_maker(dag_id=dag_id, max_active_tasks=16):
            EmptyOperator(task_id="cannot_run", pool="some_pool", pool_slots=4)
            EmptyOperator(task_id="can_run", pool="some_pool", pool_slots=1)

        scheduler_job = Job()
        self.job_runner = SchedulerJobRunner(job=scheduler_job)

Review comment (bug):
It’d probably be easier for debugging if this just do

```python
assert dag_counts == {
    "dag_1300_tasks": 4,
    "dag_1200_tasks": 4,
    ...
}
```

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

        session.add(Pool(pool="test_pool1", slots=64, description="with_slots1", include_deferred=False))
        session.add(Pool(pool="test_pool2", slots=64, description="with_slots2", include_deferred=False))
        session.flush()

        # This dag_run will have 2 RUNNING tasks and 10 SCHEDULED tasks.
        task_maker(
            dag_maker,
            session,
            dag_id="dag_12_tasks",
            task_num=12,
            max_active_tasks=4,
            running_num=2,
            run_id="run1",
        )

        queued_tis = self.job_runner._executable_task_instances_to_queued(
            max_tis=scheduler_job.executor.slots_available, session=session
        )

        assert queued_tis is not None
        # Only 2 tasks should have been queued.
        assert len(queued_tis) == 2

        dag_counts = Counter(ti.dag_id for ti in queued_tis)

        assert len(dag_counts) == 1
        assert dag_counts == {
            "dag_12_tasks": 2,
        }, "There should be only 1 dag_id with count 2 but that isn't the case"

    @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_max_active_tasks_per_dr_limit_starvation_filter_ordering(self, dag_maker, mock_executors):
        scheduler_job = Job()