[DOP-31903] Make job_dependency tests more realistic

Author: dolfinus

tests/test_server/fixtures/factories/job.py

@@ -330,92 +330,90 @@ async def jobs_with_same_parent_job(
 @pytest_asyncio.fixture
 async def job_dependency_chain(
     async_session_maker: Callable[[], AbstractAsyncContextManager[AsyncSession]],
-) -> AsyncGenerator[tuple[tuple[Job, Job, Job]], None]:
+) -> AsyncGenerator[tuple[tuple[Job, Job, Job], ...], None]:
     """
     Fixture that creates:
-    - Parent-child hierarchy: root -> middle -> leaf via parent_job_id.
-    - For each chain job - two items in job_dependency table:
-      source_root -> root and root -> target_root
-      source_middle -> middle and middle -> target_middle
-      source_leaf -> leaf and leaf -> target_leaf
-
-    Returns (job_root, job_middle, job_leaf). Root/middle/leaf are not directly connected
-    by dependency edges, only by parent-child links.
+    - Parent-child hierarchy: dag -> task -> spark via parent_job_id
+    - Job dependency edges: task1 -> task2 and task2 -> task3
+
+    There are no relations like Dag -> Dag and Spark -> Spark.
     """
     async with async_session_maker() as async_session:
         location = await create_location(async_session)
-        job_type = await create_job_type(async_session)
-        job_root = await create_job(
+
+        job_type_dag = await create_job_type(async_session, {"type": "AIRFLOW_DAG"})
+        dag1 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "job-root"},
+            job_type_id=job_type_dag.id,
+            job_kwargs={"name": "dag1"},
         )
-        job_middle = await create_job(
+        dag2 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "job-middle", "parent_job_id": job_root.id},
+            job_type_id=job_type_dag.id,
+            job_kwargs={"name": "dag2"},
         )
-        job_leaf = await create_job(
+        dag3 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "job-leaf", "parent_job_id": job_middle.id},
+            job_type_id=job_type_dag.id,
+            job_kwargs={"name": "dag3"},
         )
-        source_root = await create_job(
+
+        job_type_task = await create_job_type(async_session, {"type": "AIRFLOW_TASK"})
+        task1 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "source-root"},
+            job_type_id=job_type_task.id,
+            job_kwargs={"name": "task1", "parent_job_id": dag1.id},
         )
-        target_root = await create_job(
+        task2 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "target-root"},
+            job_type_id=job_type_task.id,
+            job_kwargs={"name": "task2", "parent_job_id": dag2.id},
         )
-        source_middle = await create_job(
+        task3 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "source-middle"},
+            job_type_id=job_type_task.id,
+            job_kwargs={"name": "task3", "parent_job_id": dag3.id},
         )
-        target_middle = await create_job(
+
+        job_type_spark = await create_job_type(async_session, {"type": "SPARK_APPLICATION"})
+        spark1 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "target-middle"},
+            job_type_id=job_type_spark.id,
+            job_kwargs={"name": "spark1", "parent_job_id": task1.id},
         )
-        source_leaf = await create_job(
+        spark2 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "source-leaf"},
+            job_type_id=job_type_spark.id,
+            job_kwargs={"name": "spark2", "parent_job_id": task2.id},
         )
-        target_leaf = await create_job(
+        spark3 = await create_job(
             async_session,
             location_id=location.id,
-            job_type_id=job_type.id,
-            job_kwargs={"name": "target-leaf"},
+            job_type_id=job_type_spark.id,
+            job_kwargs={"name": "spark3", "parent_job_id": task3.id},
         )
+
         async_session.add_all(
             [
-                JobDependency(from_job_id=source_root.id, to_job_id=job_root.id, type="DIRECT_DEPENDENCY"),
-                JobDependency(from_job_id=job_root.id, to_job_id=target_root.id, type="DIRECT_DEPENDENCY"),
-                JobDependency(from_job_id=source_middle.id, to_job_id=job_middle.id, type="DIRECT_DEPENDENCY"),
-                JobDependency(from_job_id=job_middle.id, to_job_id=target_middle.id, type="DIRECT_DEPENDENCY"),
-                JobDependency(from_job_id=source_leaf.id, to_job_id=job_leaf.id, type="DIRECT_DEPENDENCY"),
-                JobDependency(from_job_id=job_leaf.id, to_job_id=target_leaf.id, type="DIRECT_DEPENDENCY"),
+                JobDependency(from_job_id=task1.id, to_job_id=task2.id, type="DIRECT_DEPENDENCY"),
+                JobDependency(from_job_id=task2.id, to_job_id=task3.id, type="DIRECT_DEPENDENCY"),
             ],
         )
         await async_session.commit()
         async_session.expunge_all()
 
     yield (
-        (source_root, job_root, target_root),
-        (source_middle, job_middle, target_middle),
-        (source_leaf, job_leaf, target_leaf),
+        (dag1, dag2, dag3),
+        (task1, task2, task3),
+        (spark1, spark2, spark3),
     )
 
     async with async_session_maker() as async_session:

tests/test_server/test_jobs/test_job_dependencies.py

@@ -60,12 +60,11 @@ async def test_get_job_dependencies_isolated_job(
 
 async def test_get_job_dependencies_unauthorized(
     test_client: AsyncClient,
-    job_dependency_chain: tuple[tuple[Job, Job, Job], ...],
+    job: Job,
 ):
-    (_, job_middle, _) = job_dependency_chain[1]
     response = await test_client.get(
         "v1/jobs/dependencies",
-        params={"start_node_id": job_middle.id},
+        params={"start_node_id": job.id},
     )
     assert response.status_code == HTTPStatus.UNAUTHORIZED, response.json()
     assert response.json() == {
@@ -77,57 +76,54 @@ async def test_get_job_dependencies_unauthorized(
     }
 
 
-async def test_get_job_dependencies_default_request(
+async def test_get_job_dependencies_with_direction_both(
     test_client: AsyncClient,
     job_dependency_chain: tuple[tuple[Job, Job, Job], ...],
     async_session: AsyncSession,
     mocked_user: MockedUser,
 ):
     (
-        (source_root, job_root, target_root),
-        (source_middle, job_middle, target_middle),
-        (source_leaf, job_leaf, target_leaf),
+        (_, dag2, _),
+        (task1, task2, task3),
+        (_, spark2, _),
     ) = job_dependency_chain
-    all_jobs = await enrich_jobs(
+    expected_nodes = await enrich_jobs(
         [
-            job_root,
-            job_middle,
-            job_leaf,
-            source_root,
-            target_root,
-            source_middle,
-            target_middle,
-            source_leaf,
-            target_leaf,
+            dag2,
+            task1,
+            task2,
+            task3,
+            spark2,
         ],
         async_session,
     )
 
-    response = await test_client.get(
-        "v1/jobs/dependencies",
-        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"start_node_id": job_middle.id},
-    )
-    assert response.status_code == HTTPStatus.OK, response.json()
-    assert response.json() == {
-        "relations": {
-            "parents": jobs_ancestors_to_json([job_root, job_middle, job_leaf]),
-            "dependencies": [
-                {"from": {"kind": "JOB", "id": str(from_id)}, "to": {"kind": "JOB", "id": str(to_id)}, "type": type_}
-                for from_id, to_id, type_ in sorted(
-                    [
-                        (source_root.id, job_root.id, "DIRECT_DEPENDENCY"),
-                        (job_root.id, target_root.id, "DIRECT_DEPENDENCY"),
-                        (source_middle.id, job_middle.id, "DIRECT_DEPENDENCY"),
-                        (job_middle.id, target_middle.id, "DIRECT_DEPENDENCY"),
-                        (source_leaf.id, job_leaf.id, "DIRECT_DEPENDENCY"),
-                        (job_leaf.id, target_leaf.id, "DIRECT_DEPENDENCY"),
-                    ]
-                )
-            ],
-        },
-        "nodes": {"jobs": jobs_to_json(all_jobs)},
-    }
+    for start_node in [dag2, task2, spark2]:
+        response = await test_client.get(
+            "v1/jobs/dependencies",
+            headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+            params={"start_node_id": start_node.id},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+        assert response.json() == {
+            "relations": {
+                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "dependencies": [
+                    {
+                        "from": {"kind": "JOB", "id": str(from_id)},
+                        "to": {"kind": "JOB", "id": str(to_id)},
+                        "type": type_,
+                    }
+                    for from_id, to_id, type_ in sorted(
+                        [
+                            (task1.id, task2.id, "DIRECT_DEPENDENCY"),
+                            (task2.id, task3.id, "DIRECT_DEPENDENCY"),
+                        ]
+                    )
+                ],
+            },
+            "nodes": {"jobs": jobs_to_json(expected_nodes)},
+        }
 
 
 async def test_get_job_dependencies_with_direction_upstream(
@@ -136,34 +132,32 @@ async def test_get_job_dependencies_with_direction_upstream(
     async_session: AsyncSession,
     mocked_user: MockedUser,
 ):
-    (source_root, job_root, _), (source_middle, job_middle, _), (source_leaf, job_leaf, _) = job_dependency_chain
+    (_, dag2, _), (task1, task2, _), (_, spark2, _) = job_dependency_chain
     expected_nodes = await enrich_jobs(
-        [job_root, job_middle, job_leaf, source_root, source_middle, source_leaf],
+        [task1, dag2, task2, spark2],
         async_session,
     )
 
-    response = await test_client.get(
-        "v1/jobs/dependencies",
-        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"start_node_id": job_middle.id, "direction": "UPSTREAM"},
-    )
-    assert response.status_code == HTTPStatus.OK, response.json()
-    assert response.json() == {
-        "relations": {
-            "parents": jobs_ancestors_to_json([job_root, job_middle, job_leaf]),
-            "dependencies": [
-                {"from": {"kind": "JOB", "id": str(from_id)}, "to": {"kind": "JOB", "id": str(to_id)}, "type": type_}
-                for from_id, to_id, type_ in sorted(
-                    [
-                        (source_root.id, job_root.id, "DIRECT_DEPENDENCY"),
-                        (source_middle.id, job_middle.id, "DIRECT_DEPENDENCY"),
-                        (source_leaf.id, job_leaf.id, "DIRECT_DEPENDENCY"),
-                    ]
-                )
-            ],
-        },
-        "nodes": {"jobs": jobs_to_json(expected_nodes)},
-    }
+    for start_node in [dag2, task2, spark2]:
+        response = await test_client.get(
+            "v1/jobs/dependencies",
+            headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+            params={"start_node_id": start_node.id, "direction": "UPSTREAM"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+        assert response.json() == {
+            "relations": {
+                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "dependencies": [
+                    {
+                        "from": {"kind": "JOB", "id": str(task1.id)},
+                        "to": {"kind": "JOB", "id": str(task2.id)},
+                        "type": "DIRECT_DEPENDENCY",
+                    },
+                ],
+            },
+            "nodes": {"jobs": jobs_to_json(expected_nodes)},
+        }
 
 
 async def test_get_job_dependencies_with_direction_downstream(
@@ -172,31 +166,29 @@ async def test_get_job_dependencies_with_direction_downstream(
     async_session: AsyncSession,
     mocked_user: MockedUser,
 ):
-    (_, job_root, target_root), (_, job_middle, target_middle), (_, job_leaf, target_leaf) = job_dependency_chain
+    (_, dag2, _), (_, task2, task3), (_, spark2, _) = job_dependency_chain
     expected_nodes = await enrich_jobs(
-        [job_root, job_middle, job_leaf, target_root, target_middle, target_leaf],
+        [dag2, task2, spark2, task3],
         async_session,
     )
 
-    response = await test_client.get(
-        "v1/jobs/dependencies",
-        headers={"Authorization": f"Bearer {mocked_user.access_token}"},
-        params={"start_node_id": job_middle.id, "direction": "DOWNSTREAM"},
-    )
-    assert response.status_code == HTTPStatus.OK, response.json()
-    assert response.json() == {
-        "relations": {
-            "parents": jobs_ancestors_to_json([job_root, job_middle, job_leaf]),
-            "dependencies": [
-                {"from": {"kind": "JOB", "id": str(from_id)}, "to": {"kind": "JOB", "id": str(to_id)}, "type": type_}
-                for from_id, to_id, type_ in sorted(
-                    [
-                        (job_root.id, target_root.id, "DIRECT_DEPENDENCY"),
-                        (job_middle.id, target_middle.id, "DIRECT_DEPENDENCY"),
-                        (job_leaf.id, target_leaf.id, "DIRECT_DEPENDENCY"),
-                    ]
-                )
-            ],
-        },
-        "nodes": {"jobs": jobs_to_json(expected_nodes)},
-    }
+    for start_node in [dag2, task2, spark2]:
+        response = await test_client.get(
+            "v1/jobs/dependencies",
+            headers={"Authorization": f"Bearer {mocked_user.access_token}"},
+            params={"start_node_id": start_node.id, "direction": "DOWNSTREAM"},
+        )
+        assert response.status_code == HTTPStatus.OK, response.json()
+        assert response.json() == {
+            "relations": {
+                "parents": jobs_ancestors_to_json([task2, spark2]),
+                "dependencies": [
+                    {
+                        "from": {"kind": "JOB", "id": str(task2.id)},
+                        "to": {"kind": "JOB", "id": str(task3.id)},
+                        "type": "DIRECT_DEPENDENCY",
+                    },
+                ],
+            },
+            "nodes": {"jobs": jobs_to_json(expected_nodes)},
+        }