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

Code before:
        sync_dag_to_db(dag_v2, bundle_name="test", session=session)
        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 (now a simple task) and task_b

        # Find task_a - it should be a simple task now (no children)
        task_a_node = next((n for n in nodes if n["id"] == "task_a"), None)
        assert task_a_node is not None
        # After conversion from TaskGroup to task, children should be None or not present
        assert task_a_node.get("children") is None or "children" not in task_a_node

    # Tests for root, include_upstream, and include_downstream parameters
    @pytest.mark.parametrize(
        ("params", "expected_task_ids", "description"),
        [
            pytest.param(
                "root=task_c",
                ["task_c"],
                "root only returns just that task",
                id="root_only",
            ),
            pytest.param(
                "root=task_c&include_upstream=true",
                ["task_a", "task_b", "task_c"],
                "root + include_upstream returns the root task and all upstream tasks",
                id="root_upstream",
            ),
            pytest.param(
                "root=task_c&include_downstream=true",
                ["task_c", "task_d", "task_e", "task_f"],
                "root + include_downstream returns the root task and all downstream tasks including historical",
                id="root_downstream_with_historical",
            ),
        ],
    )

Review comment (suggestion):
Non-blocking nit: Could we assert with the full expected JSON result instead of partial assertions? 

Those comments are clear and we can leave them as is, but having the expected JSON response would make it even clearer than the multi-condition assertion.

Code after:
        )

        # Verify v2 structure shows TaskGroup with children
        response_v2 = test_client.get(f"/grid/structure/{dag_id}")
        assert response_v2.status_code == 200
        nodes_v2 = response_v2.json()
        assert nodes_v2 == [
            {
                "id": "task_a",
                "label": "task_a",
                "children": [
                    {"id": "task_a.task_a1", "label": "task_a1"},
                    {"id": "task_a.task_a2", "label": "task_a2"},
                ],
            },
            {"id": "task_b", "label": "task_b"},
        ]

    # Tests for root, include_upstream, and include_downstream parameters
    @pytest.mark.parametrize(
        ("params", "expected_task_ids", "description"),
        [
            pytest.param(
                "root=task_c",
                ["task_c"],
                "root only returns just that task",
                id="root_only",
            ),
            pytest.param(
                "root=task_c&include_upstream=true",
                ["task_a", "task_b", "task_c"],
                "root + include_upstream returns the root task and all upstream tasks",
                id="root_upstream",
            ),
            pytest.param(
                "root=task_c&include_downstream=true",
                ["task_c", "task_d", "task_e", "task_f"],
                "root + include_downstream returns the root task and all downstream tasks including historical",
                id="root_downstream_with_historical",
            ),
        ],
    )
    def test_structure_with_root_linear_dag(self, test_client, params, expected_task_ids, description):
        """Test root, include_upstream, and include_downstream parameters on linear DAG."""
        response = test_client.get(f"/grid/structure/{DAG_ID_5}?{params}")
        assert response.status_code == 200
        nodes = response.json()
        task_ids = sorted([node["id"] for node in nodes])
        assert task_ids == expected_task_ids, description

    # Tests for non-linear DAG structure