fix: add expected STATE messages to integration test for direct-return approach

Co-Authored-By: gl_anatolii.yatsuk <gl_anatolii.yatsuk@airbyte.io>

Author: devin-ai-integration[bot]

unit_tests/sources/test_source_read.py

@@ -6,10 +6,12 @@
 from unittest.mock import Mock
 
 import freezegun
+from airbyte_protocol_dataclasses.models import AirbyteStateType
 
 from airbyte_cdk.models import (
     AirbyteMessage,
     AirbyteRecordMessage,
+    AirbyteStateMessage,
     AirbyteStream,
     AirbyteStreamStatus,
     AirbyteStreamStatusTraceMessage,
@@ -22,11 +24,15 @@
     TraceType,
 )
 from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.models.airbyte_protocol import (
+    AirbyteStateBlob,
+    AirbyteStreamState,
+)
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
 from airbyte_cdk.sources.concurrent_source.concurrent_source_adapter import ConcurrentSourceAdapter
 from airbyte_cdk.sources.message import InMemoryMessageRepository
-from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY, Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
 from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
 from airbyte_cdk.sources.streams.core import StreamData
@@ -212,6 +218,16 @@ def test_concurrent_source_yields_the_same_messages_as_abstract_source_when_no_e
                 emitted_at=1577836800000,
             ),
         ),
+        AirbyteMessage(
+            type=MessageType.STATE,
+            state=AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_descriptor=StreamDescriptor(name="stream0"),
+                    stream_state=AirbyteStateBlob(**{NO_CURSOR_STATE_KEY: True}),
+                ),
+            ),
+        ),
         AirbyteMessage(
             type=MessageType.TRACE,
             trace=AirbyteTraceMessage(
@@ -283,6 +299,16 @@ def test_concurrent_source_yields_the_same_messages_as_abstract_source_when_no_e
                 emitted_at=1577836800000,
             ),
         ),
+        AirbyteMessage(
+            type=MessageType.STATE,
+            state=AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_descriptor=StreamDescriptor(name="stream1"),
+                    stream_state=AirbyteStateBlob(**{NO_CURSOR_STATE_KEY: True}),
+                ),
+            ),
+        ),
         AirbyteMessage(
             type=MessageType.TRACE,
             trace=AirbyteTraceMessage(