test: add missing unit tests for GroupingPartitionRouter, active_groups check, and get_partition_router

Co-Authored-By: unknown <>

Author: devin-ai-integration[bot]

unit_tests/sources/declarative/test_concurrent_declarative_source.py

@@ -5351,3 +5351,147 @@ def test_apply_stream_groups_allows_parent_child_in_different_groups():
 
     assert parent.block_simultaneous_read == "group_a"
     assert child.block_simultaneous_read == "group_b"
+
+
+def _make_child_stream_with_grouping_router(
+    child_name: str, parent_stream: DefaultStream
+) -> DefaultStream:
+    """Create a DefaultStream with GroupingPartitionRouter wrapping SubstreamPartitionRouter."""
+    from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
+        ConcurrentCursorFactory,
+        ConcurrentPerPartitionCursor,
+    )
+    from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+        GroupingPartitionRouter,
+    )
+    from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+        ParentStreamConfig,
+        SubstreamPartitionRouter,
+    )
+    from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
+        DeclarativePartitionFactory,
+        StreamSlicerPartitionGenerator,
+    )
+    from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
+    from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import (
+        EpochValueConcurrentStreamStateConverter,
+    )
+
+    substream_router = SubstreamPartitionRouter(
+        parent_stream_configs=[
+            ParentStreamConfig(
+                stream=parent_stream,
+                parent_key="id",
+                partition_field="parent_id",
+                config={},
+                parameters={},
+            )
+        ],
+        config={},
+        parameters={},
+    )
+
+    grouping_router = GroupingPartitionRouter(
+        group_size=10,
+        underlying_partition_router=substream_router,
+        config={},
+    )
+
+    cursor_factory = ConcurrentCursorFactory(lambda *args, **kwargs: Mock())
+    message_repository = InMemoryMessageRepository()
+    state_converter = EpochValueConcurrentStreamStateConverter()
+
+    per_partition_cursor = ConcurrentPerPartitionCursor(
+        cursor_factory=cursor_factory,
+        partition_router=grouping_router,
+        stream_name=child_name,
+        stream_namespace=None,
+        stream_state={},
+        message_repository=message_repository,
+        connector_state_manager=Mock(),
+        connector_state_converter=state_converter,
+        cursor_field=Mock(cursor_field_key="updated_at"),
+    )
+
+    partition_factory = Mock(spec=DeclarativePartitionFactory)
+    partition_generator = StreamSlicerPartitionGenerator(
+        partition_factory=partition_factory,
+        stream_slicer=per_partition_cursor,
+    )
+
+    cursor = FinalStateCursor(
+        stream_name=child_name, stream_namespace=None, message_repository=message_repository
+    )
+    return DefaultStream(
+        partition_generator=partition_generator,
+        name=child_name,
+        json_schema={},
+        primary_key=[],
+        cursor_field=None,
+        logger=logging.getLogger(f"test.{child_name}"),
+        cursor=cursor,
+    )
+
+
+def test_apply_stream_groups_raises_on_parent_child_in_same_group_with_grouping_router():
+    """Test _apply_stream_groups detects deadlock when GroupingPartitionRouter wraps SubstreamPartitionRouter."""
+    parent = _make_default_stream("parent_stream")
+    child = _make_child_stream_with_grouping_router("child_stream", parent)
+
+    source = Mock()
+    source._source_config = {
+        "stream_groups": {
+            "my_group": {
+                "streams": [
+                    {"name": "parent_stream", "type": "DeclarativeStream"},
+                    {"name": "child_stream", "type": "DeclarativeStream"},
+                ],
+                "action": {"type": "BlockSimultaneousSyncsAction"},
+            }
+        }
+    }
+
+    with pytest.raises(ValueError, match="child stream must not share a group with its parent"):
+        ConcurrentDeclarativeSource._apply_stream_groups(source, [parent, child])
+
+
+@pytest.mark.parametrize(
+    "stream_factory,expected_type",
+    [
+        pytest.param(
+            lambda: _make_default_stream("plain_stream"),
+            type(None),
+            id="no_partition_router_returns_none",
+        ),
+        pytest.param(
+            lambda: _make_child_stream_with_parent("child", _make_default_stream("parent")),
+            "SubstreamPartitionRouter",
+            id="substream_returns_substream_router",
+        ),
+        pytest.param(
+            lambda: _make_child_stream_with_grouping_router(
+                "child", _make_default_stream("parent")
+            ),
+            "GroupingPartitionRouter",
+            id="grouping_returns_grouping_router",
+        ),
+    ],
+)
+def test_get_partition_router(stream_factory, expected_type):
+    """Test DefaultStream.get_partition_router returns the correct router type."""
+    from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+        GroupingPartitionRouter,
+    )
+    from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+        SubstreamPartitionRouter,
+    )
+
+    stream = stream_factory()
+    router = stream.get_partition_router()
+
+    if expected_type is type(None):
+        assert router is None
+    elif expected_type == "SubstreamPartitionRouter":
+        assert isinstance(router, SubstreamPartitionRouter)
+    elif expected_type == "GroupingPartitionRouter":
+        assert isinstance(router, GroupingPartitionRouter)

unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py

@@ -28,6 +28,9 @@
 )
 from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
+    GroupingPartitionRouter,
+)
 from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
     SubstreamPartitionRouter,
 )
@@ -1440,3 +1443,85 @@ def test_is_done_raises_when_partition_generation_queue_not_empty():
 
     with pytest.raises(AirbyteTracedException, match="remained in the partition generation queue"):
         handler.is_done()
+
+
+def test_is_done_raises_when_active_groups_not_empty():
+    """Test is_done raises AirbyteTracedException if active groups remain after all streams complete."""
+    partition_enqueuer = Mock(spec=PartitionEnqueuer)
+    thread_pool_manager = Mock(spec=ThreadPoolManager)
+    logger = Mock(spec=logging.Logger)
+    slice_logger = Mock(spec=SliceLogger)
+    message_repository = Mock(spec=MessageRepository)
+    message_repository.consume_queue.return_value = []
+    partition_reader = Mock(spec=PartitionReader)
+
+    stream = Mock(spec=AbstractStream)
+    stream.name = "stuck_stream"
+    stream.block_simultaneous_read = "my_group"
+    stream.as_airbyte_stream.return_value = AirbyteStream(
+        name="stuck_stream",
+        json_schema={},
+        supported_sync_modes=[SyncMode.full_refresh],
+    )
+
+    handler = ConcurrentReadProcessor(
+        [stream],
+        partition_enqueuer,
+        thread_pool_manager,
+        logger,
+        slice_logger,
+        message_repository,
+        partition_reader,
+    )
+
+    # Mark stream as done but leave the group active (simulating a bug)
+    handler._streams_done.add("stuck_stream")
+    handler._stream_instances_to_start_partition_generation.clear()
+    handler._active_groups["my_group"] = {"stuck_stream"}
+
+    with pytest.raises(
+        AirbyteTracedException, match="still active after all streams were marked done"
+    ):
+        handler.is_done()
+
+
+def test_collect_parent_stream_names_unwraps_grouping_partition_router():
+    """Test _collect_all_parent_stream_names unwraps GroupingPartitionRouter to find parents."""
+    partition_enqueuer = Mock(spec=PartitionEnqueuer)
+    thread_pool_manager = Mock(spec=ThreadPoolManager)
+    logger = Mock(spec=logging.Logger)
+    slice_logger = Mock(spec=SliceLogger)
+    message_repository = Mock(spec=MessageRepository)
+    message_repository.consume_queue.return_value = []
+    partition_reader = Mock(spec=PartitionReader)
+
+    parent_stream = Mock(spec=AbstractStream)
+    parent_stream.name = "parent"
+    parent_stream.block_simultaneous_read = ""
+
+    # Child has a GroupingPartitionRouter wrapping a SubstreamPartitionRouter
+    child_stream = Mock(spec=DefaultStream)
+    child_stream.name = "child"
+    child_stream.block_simultaneous_read = ""
+
+    mock_substream_router = Mock(spec=SubstreamPartitionRouter)
+    mock_parent_config = Mock()
+    mock_parent_config.stream = parent_stream
+    mock_substream_router.parent_stream_configs = [mock_parent_config]
+
+    mock_grouping_router = Mock(spec=GroupingPartitionRouter)
+    mock_grouping_router.underlying_partition_router = mock_substream_router
+    child_stream.get_partition_router.return_value = mock_grouping_router
+
+    handler = ConcurrentReadProcessor(
+        [parent_stream, child_stream],
+        partition_enqueuer,
+        thread_pool_manager,
+        logger,
+        slice_logger,
+        message_repository,
+        partition_reader,
+    )
+
+    parent_names = handler._collect_all_parent_stream_names("child")
+    assert parent_names == {"parent"}