refactor(cdk): move memory check before yield, test observable behavior via caplog

Co-Authored-By: bot_apk <apk@cognition.ai>

Author: devin-ai-integration[bot]

airbyte_cdk/entrypoint.py

@@ -281,8 +281,8 @@ def read(
         stream_message_counter: DefaultDict[HashableStreamDescriptor, float] = defaultdict(float)
         try:
             for message in self.source.read(self.logger, config, catalog, state):
-                yield self.handle_record_counts(message, stream_message_counter)
                 self._memory_monitor.check_memory_usage()
+                yield self.handle_record_counts(message, stream_message_counter)
         finally:
             for message in self._emit_queued_messages(self.source):
                 yield self.handle_record_counts(message, stream_message_counter)

unit_tests/test_entrypoint.py

@@ -908,7 +908,8 @@ def _raise_on_second_call() -> None:
         for msg in entrypoint.run(parsed_args):
             messages.append(msg)
 
-    # 1. The first record was yielded before the exception
+    # 1. Only the first record was yielded; the second triggered the exception
+    #    before its yield so it is handled in the finally block instead.
     record_messages = [m for m in messages if "RECORD" in m]
     assert len(record_messages) >= 1, (
         "At least the first record should be yielded before MemoryLimitExceeded"
@@ -921,10 +922,11 @@ def _raise_on_second_call() -> None:
     )
 
     # 3. The flushed state has sourceStats.recordCount set by handle_record_counts.
-    #    Both records are yielded (and counted) before the second check_memory_usage
-    #    raises, so the counter is 2.0 at flush time.
+    #    Only the first record is yielded (and counted) because the memory check
+    #    now runs *before* yield — the second check raises before the second
+    #    record is handled, so the counter is 1.0 at flush time.
     state_json = orjson.loads(state_messages[0])
-    assert state_json["state"]["sourceStats"]["recordCount"] == 2.0
+    assert state_json["state"]["sourceStats"]["recordCount"] == 1.0
 
 
 def test_given_serialization_error_using_orjson_then_fallback_on_json(