openlineage: add unit test for listener hooks on dag run state changes and task instance failure

tests/dags/test_openlineage_execution.py

@@ -27,13 +27,16 @@
 
 
 class OpenLineageExecutionOperator(BaseOperator):
-    def __init__(self, *, stall_amount=0, **kwargs) -> None:
+    def __init__(self, *, stall_amount=0, fail=False, **kwargs) -> None:
         super().__init__(**kwargs)
         self.stall_amount = stall_amount
+        self.fail = fail
 
     def execute(self, context):
         self.log.error("STALL AMOUNT %s", self.stall_amount)
         time.sleep(1)
+        if self.fail:
+            raise Exception("Failed")
 
     def get_openlineage_facets_on_start(self):
         return OperatorLineage(inputs=[Dataset(namespace="test", name="on-start")])
@@ -43,6 +46,11 @@ def get_openlineage_facets_on_complete(self, task_instance):
         time.sleep(self.stall_amount)
         return OperatorLineage(inputs=[Dataset(namespace="test", name="on-complete")])
 
+    def get_openlineage_facets_on_failure(self, task_instance):
+        self.log.error("STALL AMOUNT %s", self.stall_amount)
+        time.sleep(self.stall_amount)
+        return OperatorLineage(inputs=[Dataset(namespace="test", name="on-failure")])
+
 
 with DAG(
     dag_id="test_openlineage_execution",
@@ -57,3 +65,5 @@ def get_openlineage_facets_on_complete(self, task_instance):
     mid_stall = OpenLineageExecutionOperator(task_id="execute_mid_stall", stall_amount=15)
 
     long_stall = OpenLineageExecutionOperator(task_id="execute_long_stall", stall_amount=30)
+
+    fail = OpenLineageExecutionOperator(task_id="execute_fail", fail=True)

tests/providers/openlineage/plugins/test_execution.py

@@ -124,6 +124,17 @@ def test_not_stalled_task_emits_proper_lineage(self):
             assert has_value_in_events(events, ["inputs", "name"], "on-start")
             assert has_value_in_events(events, ["inputs", "name"], "on-complete")
 
+        @pytest.mark.db_test
+        @conf_vars({("openlineage", "transport"): f'{{"type": "file", "log_file_path": "{listener_path}"}}'})
+        def test_not_stalled_failing_task_emits_proper_lineage(self):
+            task_name = "execute_fail"
+            run_id = "test_failure"
+            self.setup_job(task_name, run_id)
+
+            events = get_sorted_events(tmp_dir)
+            assert has_value_in_events(events, ["inputs", "name"], "on-start")
+            assert has_value_in_events(events, ["inputs", "name"], "on-failure")
+
         @conf_vars(
             {
                 ("openlineage", "transport"): f'{{"type": "file", "log_file_path": "{listener_path}"}}',

tests/providers/openlineage/plugins/test_listener.py

@@ -606,6 +606,49 @@ def test_listener_on_dag_run_state_changes_configure_process_pool_size(mock_exec
     mock_executor.return_value.submit.assert_called_once()
 
 
+class MockExecutor:
+    def __init__(self, *args, **kwargs):
+        self.submitted = False
+        self.succeeded = False
+        self.result = None
+
+    def submit(self, fn, /, *args, **kwargs):
+        self.submitted = True
+        try:
+            fn(*args, **kwargs)
+            self.succeeded = True
+        except Exception:
+            pass
+        return MagicMock()
+
+    def shutdown(self, *args, **kwargs):
+        print("Shutting down")
+
+
+@pytest.mark.parametrize(
+    ("method", "dag_run_state"),
+    [
+        ("on_dag_run_running", DagRunState.RUNNING),
+        ("on_dag_run_success", DagRunState.SUCCESS),
+        ("on_dag_run_failed", DagRunState.FAILED),
+    ],
+)
+@patch("airflow.providers.openlineage.plugins.adapter.OpenLineageAdapter.emit")
+def test_listener_on_dag_run_state_changes(mock_emit, method, dag_run_state, create_task_instance):
+    mock_executor = MockExecutor()
+    ti = create_task_instance(dag_id="dag", task_id="op")
+    # Change the state explicitly to set end_date following the logic in the method
+    ti.dag_run.set_state(dag_run_state)
+    with mock.patch(
+        "airflow.providers.openlineage.plugins.listener.ProcessPoolExecutor", return_value=mock_executor
+    ):
+        listener = OpenLineageListener()
+        getattr(listener, method)(ti.dag_run, None)
+        assert mock_executor.submitted is True
+        assert mock_executor.succeeded is True
+        mock_emit.assert_called_once()
+
+
 def test_listener_logs_failed_serialization():
     listener = OpenLineageListener()
     callback_future = Future()