Repository: apache/airflow
File: airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py (Python)

Code before:
            task_b = PythonOperator(task_id="task_b", python_callable=lambda: True)

        # Create a DagRun with the old version
        dag_bag = DBDagBag()
        dag_bag.bag_dag(dag_v1, root_dag=dag_v1)
        session.commit()

        dr1 = DagRun(
            dag_id=dag_id,
            run_id="test_run_1",
            run_type=DagRunType.MANUAL,
            run_after=pendulum.datetime(2024, 1, 1, tz="UTC"),
        )
        session.add(dr1)
        session.commit()

        # Version 2: task_a becomes a TaskGroup with subtasks
        with DAG(
            dag_id=dag_id,
            start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
            schedule=None,
        ) as dag_v2:
            with TaskGroup(group_id="task_a") as task_a_group:  # noqa: F841
                task_a1 = PythonOperator(task_id="task_a1", python_callable=lambda: True)  # noqa: F841
                task_a2 = PythonOperator(task_id="task_a2", python_callable=lambda: True)  # noqa: F841
            task_b = PythonOperator(task_id="task_b", python_callable=lambda: True)  # noqa: F841

        # Update the DAG with the new version
        dag_bag.bag_dag(dag_v2, root_dag=dag_v2)
        session.commit()

        # Create another DagRun with the new version
        dr2 = DagRun(
            dag_id=dag_id,
            run_id="test_run_2",
            run_type=DagRunType.MANUAL,
            run_after=pendulum.datetime(2024, 1, 2, tz="UTC"),
        )
        session.add(dr2)
        session.commit()

        # This should not crash with TypeError: 'NoneType' object is not iterable
        response = test_client.get(f"/grid/structure/{dag_id}")
        assert response.status_code == 200

        # Verify the structure includes both versions merged correctly
        nodes = response.json()
        assert len(nodes) == 2  # task_a (group) and task_b

        # Find task_a - it should be a group now
        task_a_node = next((n for n in nodes if n["id"] == "task_a"), None)

Review comment (suggestion):
Remove the variable assignment it appears you don't need it. So you don't have to `# noqa: F841`

Code after:

        Regression test for https://github.com/apache/airflow/issues/61208
        """

        dag_id = "test_task_to_group_conversion"

        # Version 1: task_a is a simple task
        with dag_maker(
            dag_id=dag_id,
            start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
            schedule=None,
        ):
            PythonOperator(task_id="task_a", python_callable=lambda: True)
            PythonOperator(task_id="task_b", python_callable=lambda: True)

        # Create another DagRun with the new version
        dag_maker.create_dagrun(
            run_id="test_run_1",
            run_type=DagRunType.MANUAL,
            logical_date=pendulum.datetime(2024, 1, 3, tz="UTC"),
        )

        response_v1 = test_client.get(f"/grid/structure/{dag_id}")
        assert response_v1.status_code == 200
        nodes_v1 = response_v1.json()
        assert nodes_v1 == [
            {"id": "task_a", "label": "task_a"},
            {"id": "task_b", "label": "task_b"},
        ]

        # Version 2: task_a is a TaskGroup with subtasks
        with dag_maker(
            dag_id=dag_id,
            start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
            schedule=None,
            serialized=True,
        ):
            with TaskGroup(group_id="task_a"):
                PythonOperator(task_id="task_a1", python_callable=lambda: True)
                PythonOperator(task_id="task_a2", python_callable=lambda: True)
            PythonOperator(task_id="task_b", python_callable=lambda: True)

        dag_maker.create_dagrun(
            run_id="test_run_2",
            run_type=DagRunType.MANUAL,
            logical_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
        )

        # Verify v2 structure shows TaskGroup with children
        response_v2 = test_client.get(f"/grid/structure/{dag_id}")
        assert response_v2.status_code == 200