adding lots of logging to diagnose flaky test

Author: brianjlai

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 import logging
+import os
 from typing import Dict, Iterable, List, Optional, Set
 
 from airbyte_cdk.exception_handler import generate_failed_streams_error_message
@@ -154,6 +155,9 @@ def on_record(self, record: Record) -> Iterable[AirbyteMessage]:
                 )
             self._record_counter[stream.name] += 1
             stream.cursor.observe(record)
+        test_env = os.getenv("PYTEST_CURRENT_TEST")
+        if test_env and "test_concurrent_declarative_source.py" in test_env:
+            self._logger.info(f"Processing and emitting: {message.__dict__}")
         yield message
         yield from self._message_repository.consume_queue()
 

airbyte_cdk/sources/message/concurrent_repository.py

@@ -1,12 +1,16 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-
+import logging
+import os
 from queue import Queue
 from typing import Callable, Iterable
 
 from airbyte_cdk.models import AirbyteMessage, Level
+from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.message.repository import LogMessage, MessageRepository
 from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
 
+logger = logging.getLogger("airbyte")
+
 
 class ConcurrentMessageRepository(MessageRepository):
     """
@@ -25,14 +29,25 @@ def __init__(self, queue: Queue[QueueItem], message_repository: MessageRepositor
         self._queue = queue
         self._decorated_message_repository = message_repository
 
+        test_env = os.getenv("PYTEST_CURRENT_TEST")
+        self._log_messages_for_testing = (
+            test_env and "test_concurrent_declarative_source.py" in test_env
+        )
+
     def emit_message(self, message: AirbyteMessage) -> None:
+        if self._log_messages_for_testing:
+            self._log_message(message)
         self._decorated_message_repository.emit_message(message)
         for message in self._decorated_message_repository.consume_queue():
+            if self._log_messages_for_testing:
+                self._log_message(message)
             self._queue.put(message)
 
     def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
         self._decorated_message_repository.log_message(level, message_provider)
         for message in self._decorated_message_repository.consume_queue():
+            if self._log_messages_for_testing:
+                self._log_message(message)
             self._queue.put(message)
 
     def consume_queue(self) -> Iterable[AirbyteMessage]:
@@ -41,3 +56,16 @@ def consume_queue(self) -> Iterable[AirbyteMessage]:
         loading messages onto the queue processed on the main thread.
         """
         yield from []
+
+    @staticmethod
+    def _log_message(message: AirbyteMessage) -> None:
+        if message.type == MessageType.STATE:
+            if message.state and message.state.stream:
+                state = message.state.stream.stream_state.__dict__
+                logger.info(
+                    f"Processing and emitting message of type {message.type} with contents: {message.state.stream.stream_state.__dict__}"
+                )
+        else:
+            logger.info(
+                f"Processing and emitting message of type {message.type} with contents: {message.__dict__}"
+            )

unit_tests/sources/declarative/test_concurrent_declarative_source.py

@@ -1507,6 +1507,8 @@ def test_read_concurrent_with_failing_partition_in_the_middle():
             ):
                 messages.append(message)
         except AirbyteTracedException:
+            locations_states = get_states_for_stream(stream_name="locations", messages=messages)
+            assert len(locations_states) == 3
             assert (
                 get_states_for_stream(stream_name="locations", messages=messages)[
                     -1