[DOP-22750] add unit tests for 'airflow' owner

Author: TiGrib

tests/test_consumer/test_extractors/test_extractors_batch_airflow.py

@@ -80,6 +80,46 @@ def airflow_dag_run_event_start() -> OpenLineageRunEvent:
     )
 
 
+@pytest.fixture
+def airflow_dag_run_event_start_owner_airflow() -> OpenLineageRunEvent:
+    event_time = datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc)
+    run_id = UUID("01908223-0782-79b8-9495-b1c38aaee839")
+    return OpenLineageRunEvent(
+        eventType=OpenLineageRunEventType.START,
+        eventTime=event_time,
+        job=OpenLineageJob(
+            namespace="http://airflow-host:8081",
+            name="mydag",
+            facets=OpenLineageJobFacets(
+                jobType=OpenLineageJobTypeJobFacet(
+                    processingType=None,
+                    integration=OpenLineageJobIntegrationType.AIRFLOW,
+                    jobType=OpenLineageJobType.DAG,
+                ),
+            ),
+        ),
+        run=OpenLineageRun(
+            runId=run_id,
+            facets=OpenLineageRunFacets(
+                processing_engine=OpenLineageProcessingEngineRunFacet(
+                    version=Version("2.9.2"),
+                    name=OpenLineageProcessingEngineName.AIRFLOW,
+                    openlineageAdapterVersion=Version("1.10.0"),
+                ),
+                airflowDagRun=OpenLineageAirflowDagRunFacet(
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
+                    dagRun=OpenLineageAirflowDagRunInfo(
+                        run_id="manual__2024-07-05T09:04:13:979349+00:00",
+                        run_type=OpenLineageAirflowDagRunType.MANUAL,
+                        data_interval_start=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+                        data_interval_end=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+
 @pytest.fixture
 def airflow_dag_run_event_stop() -> OpenLineageRunEvent:
     event_time = datetime(2024, 7, 5, 9, 8, 5, 691973, tzinfo=timezone.utc)
@@ -160,6 +200,55 @@ def airflow_task_run_event_start() -> OpenLineageRunEvent:
     )
 
 
+@pytest.fixture
+def airflow_task_run_event_start_owner_airflow() -> OpenLineageRunEvent:
+    event_time = datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc)
+    run_id = UUID("01908223-0782-7fc0-9d69-b1df9dac2c60")
+    return OpenLineageRunEvent(
+        eventType=OpenLineageRunEventType.START,
+        eventTime=event_time,
+        job=OpenLineageJob(
+            namespace="http://airflow-host:8081",
+            name="mydag.mytask",
+            facets=OpenLineageJobFacets(
+                jobType=OpenLineageJobTypeJobFacet(
+                    processingType=None,
+                    integration=OpenLineageJobIntegrationType.AIRFLOW,
+                    jobType=OpenLineageJobType.TASK,
+                ),
+            ),
+        ),
+        run=OpenLineageRun(
+            runId=run_id,
+            facets=OpenLineageRunFacets(
+                processing_engine=OpenLineageProcessingEngineRunFacet(
+                    version=Version("2.9.2"),
+                    name=OpenLineageProcessingEngineName.AIRFLOW,
+                    openlineageAdapterVersion=Version("1.10.0"),
+                ),
+                airflow=OpenLineageAirflowTaskRunFacet(
+                    dag=OpenLineageAirflowDagInfo(dag_id="mydag", owner="airflow"),
+                    dagRun=OpenLineageAirflowDagRunInfo(
+                        run_id="manual__2024-07-05T09:04:13:979349+00:00",
+                        run_type=OpenLineageAirflowDagRunType.MANUAL,
+                        data_interval_start=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+                        data_interval_end=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+                    ),
+                    task=OpenLineageAirflowTaskInfo(
+                        task_id="mytask",
+                    ),
+                    taskInstance=OpenLineageAirflowTaskInstanceInfo(
+                        try_number=1,
+                        log_url=(
+                            "http://airflow-host:8081/dags/mydag/grid?tab=logs&dag_run_id=manual__2024-07-05T09%3A04%3A13%3A979349%2B00%3A00&task_id=mytask"
+                        ),
+                    ),
+                ),
+            ),
+        ),
+    )
+
+
 @pytest.fixture
 def airflow_task_run_event_stop() -> OpenLineageRunEvent:
     event_time = datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc)
@@ -222,6 +311,41 @@ def extracted_airflow_task_job(
     )
 
 
+@pytest.fixture
+def extracted_airflow_dag_run_owner_airflow(
+    extracted_airflow_dag_job: JobDTO,
+) -> RunDTO:
+    return RunDTO(
+        id=UUID("01908223-0782-79b8-9495-b1c38aaee839"),
+        job=extracted_airflow_dag_job,
+        status=RunStatusDTO.SUCCEEDED,
+        started_at=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+        start_reason=RunStartReasonDTO.MANUAL,
+        ended_at=datetime(2024, 7, 5, 9, 8, 5, 691973, tzinfo=timezone.utc),
+        external_id="manual__2024-07-05T09:04:13:979349+00:00",
+        persistent_log_url="http://airflow-host:8081/dags/mydag/grid?dag_run_id=manual__2024-07-05T09%3A04%3A13%3A979349%2B00%3A00",
+    )
+
+
+@pytest.fixture
+def extracted_airflow_task_run_owner_airflow(
+    extracted_airflow_task_job: JobDTO,
+) -> RunDTO:
+    return RunDTO(
+        id=UUID("01908223-0782-7fc0-9d69-b1df9dac2c60"),
+        job=extracted_airflow_task_job,
+        status=RunStatusDTO.SUCCEEDED,
+        started_at=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+        start_reason=RunStartReasonDTO.MANUAL,
+        ended_at=datetime(2024, 7, 5, 9, 4, 13, 979349, tzinfo=timezone.utc),
+        external_id="manual__2024-07-05T09:04:13:979349+00:00",
+        attempt="1",
+        persistent_log_url=(
+            "http://airflow-host:8081/dags/mydag/grid?tab=logs&dag_run_id=manual__2024-07-05T09%3A04%3A13%3A979349%2B00%3A00&task_id=mytask"
+        ),
+    )
+
+
 @pytest.fixture
 def extracted_airflow_dag_run(
     extracted_airflow_dag_job: JobDTO,
@@ -313,3 +437,54 @@ def test_extractors_extract_batch_airflow(
     assert not extracted.operations()
     assert not extracted.inputs()
     assert not extracted.outputs()
+
+
+@pytest.mark.parametrize(
+    "input_transformation",
+    [
+        # receiving data out of order does not change result
+        pytest.param(
+            list,
+            id="preserve order",
+        ),
+        pytest.param(
+            reversed,
+            id="reverse order",
+        ),
+    ],
+)
+def test_extractors_extract_batch_airflow_with_airflow_owner(
+    airflow_dag_run_event_start_owner_airflow: OpenLineageRunEvent,
+    airflow_dag_run_event_stop: OpenLineageRunEvent,
+    airflow_task_run_event_start_owner_airflow: OpenLineageRunEvent,
+    airflow_task_run_event_stop: OpenLineageRunEvent,
+    extracted_airflow_location: LocationDTO,
+    extracted_airflow_dag_job: JobDTO,
+    extracted_airflow_task_job: JobDTO,
+    extracted_airflow_dag_run_owner_airflow: RunDTO,
+    extracted_airflow_task_run_owner_airflow: RunDTO,
+    input_transformation,
+):
+    events = [
+        airflow_dag_run_event_start_owner_airflow,
+        airflow_task_run_event_start_owner_airflow,
+        airflow_task_run_event_stop,
+        airflow_dag_run_event_stop,
+    ]
+
+    extracted = extract_batch(input_transformation(events))
+
+    assert extracted.locations() == [extracted_airflow_location]
+    assert extracted.jobs() == [extracted_airflow_dag_job, extracted_airflow_task_job]
+    assert extracted.runs() == [
+        extracted_airflow_dag_run_owner_airflow,
+        extracted_airflow_task_run_owner_airflow,
+    ]
+
+    assert not extracted.datasets()
+    assert not extracted.dataset_symlinks()
+    assert not extracted.schemas()
+    assert not extracted.users()
+    assert not extracted.operations()
+    assert not extracted.inputs()
+    assert not extracted.outputs()

tests/test_consumer/test_handlers/test_runs_handler_airflow.py

@@ -77,6 +77,10 @@ async def test_runs_handler_airflow(
     runs = run_scalars.all()
     assert len(runs) == 2
 
+    user_query = select(User).where(User.name == "myuser")
+    user_scalars = await async_session.scalars(user_query)
+    user = user_scalars.one_or_none()
+
     dag_run = runs[0]
     assert dag_run.id == UUID("01908223-0782-79b8-9495-b1c38aaee839")
     assert dag_run.created_at == datetime(2024, 7, 5, 9, 4, 12, 162000, tzinfo=timezone.utc)
@@ -98,9 +102,7 @@ async def test_runs_handler_airflow(
     assert task_run.status == RunStatus.SUCCEEDED
     assert task_run.started_at == datetime(2024, 7, 5, 9, 4, 20, 783845, tzinfo=timezone.utc)
     assert task_run.ended_at == datetime(2024, 7, 5, 9, 7, 37, 858423, tzinfo=timezone.utc)
-    user_query = select(User).where(User.name == "myuser")
-    user_scalars = await async_session.scalars(user_query)
-    user = user_scalars.one_or_none()
+    assert task_run.started_by_user_id is not None
     assert task_run.started_by_user_id == user.id
     assert task_run.start_reason == RunStartReason.MANUAL
     assert task_run.end_reason is None