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

Code before:
        run_id = DagRun.generate_run_id(
            run_type=run_type,
            logical_date=logical_date,
            run_after=run_after,
        )
        return dag.create_dagrun(
            run_id=run_id,
            logical_date=logical_date,
            data_interval=data_interval,
            run_after=run_after,
            run_type=run_type,
            state=state,
            start_date=start_date,
            triggered_by=DagRunTriggeredByType.TEST,
        )

    return _create_dagrun


def task_maker(
    dag_maker, session, dag_id: str, task_num: int, max_active_tasks: int, run_id: str | None = None
):
    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:  # 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}")

    # 'create_dagrun' uses a kwargs dict to check whether parameters
    # are present. Do the same for simplicity.
    # 'logical_date' is used to create the 'run_id'. Set it to 'now',
    # in order to get distinct run ids, if a value isn't provided.
    kwargs = {
        "run_type": DagRunType.SCHEDULED,
        "logical_date": timezone.utcnow(),
    }

    if run_id is not None:
        kwargs["run_id"] = run_id

    dag_run = dag_maker.create_dagrun(**kwargs)

    task_tis = {}

    tis_list = []

Review comment (suggestion):
Couple test cases worth adding:

1. **Starvation filter ordering**: dag run with tasks in mixed pools (some starved, some not). Verify non-starved pool tasks aren't excluded because starved-pool tasks consumed row_number slots.

2. **Partial capacity**: dag run with `max_active_tasks=4` where 2 are already RUNNING + 10 SCHEDULED. Verify query returns only 2 (not 4) for that run.

Code after:
        data_interval: DataInterval,
        run_type: DagRunType,
        state: DagRunState = DagRunState.RUNNING,
        start_date: datetime.datetime | None = None,
    ) -> DagRun:
        run_after = logical_date or timezone.utcnow()
        run_id = DagRun.generate_run_id(
            run_type=run_type,
            logical_date=logical_date,
            run_after=run_after,
        )
        return dag.create_dagrun(
            run_id=run_id,
            logical_date=logical_date,
            data_interval=data_interval,
            run_after=run_after,
            run_type=run_type,
            state=state,
            start_date=start_date,
            triggered_by=DagRunTriggeredByType.TEST,
        )

    return _create_dagrun


def task_maker(
    dag_maker,
    session,
    dag_id: str,
    task_num: int,
    max_active_tasks: int,
    running_num: int,
    run_id: str | None = None,
):
    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:  # 0, 10, 20, 30, 40, 50, ...
                weight = int(i / 2)
                dag_tasks[f"op{i}"] = EmptyOperator(
                    task_id=f"dummy{i}", priority_weight=weight, pool="test_pool1"
                )
            else:
                # No executor specified, runs on default executor
                dag_tasks[f"op{i}"] = EmptyOperator(task_id=f"dummy{i}", pool="test_pool2")

    # 'create_dagrun' uses a kwargs dict to check whether parameters
    # are present. Do the same for simplicity.
    # 'logical_date' is used to create the 'run_id'. Set it to 'now',