(airflow) 01/01: test: make flaky OL test wait for future callback done

[mo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mobuchowski pushed a commit to branch make-sure-log-callback-done
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit 17720715cce7b21a5171935ef85b7899302e654f
Author: Maciej Obuchowski <ob...@gmail.com>
AuthorDate: Thu Sep 5 13:02:55 2024 +0200

    test: make flaky OL test wait for future callback done
    
    Signed-off-by: Maciej Obuchowski <ob...@gmail.com>
---
 tests/providers/openlineage/plugins/test_listener.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/providers/openlineage/plugins/test_listener.py b/tests/providers/openlineage/plugins/test_listener.py
index 2fbe27f23d..44c204e7af 100644
--- a/tests/providers/openlineage/plugins/test_listener.py
+++ b/tests/providers/openlineage/plugins/test_listener.py
@@ -17,8 +17,8 @@
 from __future__ import annotations
 
 import datetime as dt
-import threading
 import uuid
+from concurrent.futures import Future
 from contextlib import suppress
 from typing import Callable
 from unittest import mock
@@ -606,26 +606,28 @@ def test_listener_on_dag_run_state_changes_configure_process_pool_size(mock_exec
     mock_executor.return_value.submit.assert_called_once()
 
 
-@pytest.mark.flaky(reruns=5)
 def test_listener_logs_failed_serialization():
     listener = OpenLineageListener()
-    listener.log = MagicMock()
+    callback_future = Future()
+
+    listener.log = MagicMock(side_effect=lambda: callback_future.set_result(True))
     listener.adapter = OpenLineageAdapter(
         client=OpenLineageClient(transport=ConsoleTransport(config=ConsoleConfig()))
     )
     event_time = dt.datetime.now()
-
     fut = listener.submit_callable(
         listener.adapter.dag_failed,
         dag_id="",
         run_id="",
         end_date=event_time,
-        execution_date=threading.Thread(),
+        execution_date=callback_future,
         dag_run_state=DagRunState.FAILED,
         task_ids=["task_id"],
         msg="",
     )
     assert fut.exception(10)
+    assert callback_future.done()
+    listener.log.debug.assert_not_called()
     listener.log.warning.assert_called_once()