Fix compact tests

Author: jason810496

devel-common/src/tests_common/test_utils/file_task_handler.py

@@ -24,10 +24,10 @@
 
 import pendulum
 
-from airflow.utils.log.file_task_handler import StructuredLogMessage
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
-    from airflow.utils.log.file_task_handler import ParsedLog
+    from airflow.utils.log.file_task_handler import ParsedLog, StructuredLogMessage
 
 
 def events(logs: Iterable[StructuredLogMessage], skip_source_info=True) -> list[str]:
@@ -60,6 +60,9 @@ def mock_parsed_logs_factory(
     Create a list of ParsedLog objects with the specified start datetime and count.
     Each ParsedLog object contains a timestamp and a list of StructuredLogMessage objects.
     """
+    if AIRFLOW_V_3_0_PLUS:
+        from airflow.utils.log.file_task_handler import StructuredLogMessage
+
     parsed_logs: list[ParsedLog] = []
     for i in range(count):
         timestamp: datetime = start_datetime + pendulum.duration(seconds=i)

providers/amazon/tests/unit/amazon/aws/log/test_cloudwatch_task_handler.py

@@ -260,10 +260,8 @@ def test_read(self, monkeypatch):
             ],
         )
         if AIRFLOW_V_2_10_PLUS:
-            monkeypatch.setattr(
-                self.cloudwatch_task_handler,
-                "_read_from_logs_server",
-                lambda ti, worker_log_rel_path, log_metadata: ([], []),
+            self.cloudwatch_task_handler._read_from_logs_server = mock.Mock(
+                return_value=([], []),
             )
             msg_template = textwrap.dedent("""
                  INFO - ::group::Log message source details

providers/amazon/tests/unit/amazon/aws/log/test_s3_task_handler.py

@@ -261,7 +261,7 @@ def test_read(self):
             actual = log[0][0][-1]
             assert f"*** Found logs in s3:\n***   * {expected_s3_uri}\n" in actual
             assert actual.endswith("Line 4")
-            assert metadata == [{"end_of_log": True, "first_time_read": 33}]
+            assert metadata == [{"end_of_log": True, "log_pos": 33}]
 
     def test_read_when_s3_log_missing(self):
         ti = copy.copy(self.ti)

providers/celery/tests/unit/celery/log_handlers/test_log_handlers.py

@@ -81,10 +81,16 @@ def test__read_for_celery_executor_fallbacks_to_worker(self, create_task_instanc
             fth = FileTaskHandler("")
 
             fth._read_from_logs_server = mock.Mock()
-            fth._read_from_logs_server.return_value = (
-                ["this message"],
-                [convert_list_to_stream(["this", "log", "content"])],
-            )
+
+            # compact with 2.x and 3.x
+            if AIRFLOW_V_3_0_PLUS:
+                fth._read_from_logs_server.return_value = (
+                    ["this message"],
+                    [convert_list_to_stream(["this", "log", "content"])],
+                )
+            else:
+                fth._read_from_logs_server.return_value = ["this message"], ["this\nlog\ncontent"]
+
             logs, metadata = fth._read(ti=ti, try_number=1)
             fth._read_from_logs_server.assert_called_once()
 

providers/elasticsearch/src/airflow/providers/elasticsearch/log/es_task_handler.py

@@ -55,11 +55,18 @@
     from datetime import datetime
 
     from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
+    from airflow.typing_compat import TypeAlias
 
 if AIRFLOW_V_3_0_PLUS:
-    from airflow.utils.log.file_task_handler import LogHandlerOutputStream
+    from collections.abc import Generator
+    from itertools import chain
+    from typing import Union
 
-    EsLogMsgType = LogHandlerOutputStream
+    from airflow.utils.log.file_task_handler import StructuredLogMessage
+
+    StructuredLogStream: TypeAlias = Generator[StructuredLogMessage, None, None]
+
+    EsLogMsgType = Union[StructuredLogStream, chain[StructuredLogMessage]]
 else:
     EsLogMsgType = list[tuple[str, str]]  # type: ignore[misc]
 

providers/opensearch/src/airflow/providers/opensearch/log/os_task_handler.py

@@ -22,6 +22,7 @@
 import sys
 import time
 from collections import defaultdict
+from collections.abc import Generator
 from datetime import datetime
 from operator import attrgetter
 from typing import TYPE_CHECKING, Any, Callable, Literal
@@ -44,12 +45,19 @@
 
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
+    from airflow.typing_compat import TypeAlias
 
 
 if AIRFLOW_V_3_0_PLUS:
-    from airflow.utils.log.file_task_handler import LogHandlerOutputStream
+    from collections.abc import Generator
+    from itertools import chain
+    from typing import Union
 
-    OsLogMsgType = LogHandlerOutputStream
+    from airflow.utils.log.file_task_handler import StructuredLogMessage
+
+    StructuredLogStream: TypeAlias = Generator[StructuredLogMessage, None, None]
+
+    OsLogMsgType = Union[StructuredLogStream, chain[StructuredLogMessage]]
 else:
     OsLogMsgType = list[tuple[str, str]]  # type: ignore[misc]