#
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
#
import logging
import unittest
from unittest.mock import Mock, call

import freezegun
import pytest

from airbyte_cdk.models import (
    AirbyteLogMessage,
    AirbyteMessage,
    AirbyteRecordMessage,
    AirbyteStream,
    AirbyteStreamStatus,
    AirbyteStreamStatusTraceMessage,
    AirbyteTraceMessage,
    FailureType,
    StreamDescriptor,
    SyncMode,
    TraceType,
)
from airbyte_cdk.models import Level as LogLevel
from airbyte_cdk.models import Type as MessageType
from airbyte_cdk.sources.concurrent_source.concurrent_read_processor import ConcurrentReadProcessor
from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
    PartitionGenerationCompletedSentinel,
)
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,
)
from airbyte_cdk.sources.message import LogMessage, MessageRepository
from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionEnqueuer
from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader
from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel
from airbyte_cdk.sources.types import Record
from airbyte_cdk.sources.utils.slice_logger import SliceLogger
from airbyte_cdk.utils.traced_exception import AirbyteTracedException

_STREAM_NAME = "stream"
_ANOTHER_STREAM_NAME = "stream2"
_ANY_AIRBYTE_MESSAGE = Mock(spec=AirbyteMessage)
_IS_SUCCESSFUL = True


class TestConcurrentReadProcessor(unittest.TestCase):
    def setUp(self):
        self._partition_enqueuer = Mock(spec=PartitionEnqueuer)
        self._thread_pool_manager = Mock(spec=ThreadPoolManager)

        self._an_open_partition = Mock(spec=Partition)
        self._log_message = Mock(spec=LogMessage)
        self._an_open_partition.to_slice.return_value = self._log_message
        self._an_open_partition.stream_name.return_value = _STREAM_NAME

        self._a_closed_partition = Mock(spec=Partition)
        self._a_closed_partition.stream_name.return_value = _ANOTHER_STREAM_NAME

        self._logger = Mock(spec=logging.Logger)
        self._slice_logger = Mock(spec=SliceLogger)
        self._slice_logger.create_slice_log_message.return_value = self._log_message
        self._message_repository = Mock(spec=MessageRepository)
        self._message_repository.consume_queue.return_value = []
        self._partition_reader = Mock(spec=PartitionReader)

        self._stream = Mock(spec=AbstractStream)
        self._stream.name = _STREAM_NAME
        self._stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )
        self._another_stream = Mock(spec=AbstractStream)
        self._another_stream.name = _ANOTHER_STREAM_NAME
        self._another_stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_ANOTHER_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )

        self._record_data = {"id": 1, "value": "A"}
        self._partition = Mock(spec=Partition)
        self._partition.stream_name = lambda: _STREAM_NAME
        self._record = Mock(spec=Record)
        self._record.partition = self._partition
        self._record.data = self._record_data
        self._record.stream_name = _STREAM_NAME
        self._record.file_reference = None

    def test_stream_is_not_done_initially(self):
        stream_instances_to_read_from = [self._stream]
        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        assert not handler._is_stream_done(self._stream.name)

    def test_handle_partition_done_no_other_streams_to_generate_partitions_for(self):
        stream_instances_to_read_from = [self._stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)

        sentinel = PartitionGenerationCompletedSentinel(self._stream)
        messages = list(handler.on_partition_generation_completed(sentinel))

        expected_messages = []
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_handle_last_stream_partition_done(self):
        in_order_validation_mock = Mock()
        in_order_validation_mock.attach_mock(self._another_stream, "_another_stream")
        in_order_validation_mock.attach_mock(self._message_repository, "_message_repository")
        self._message_repository.consume_queue.return_value = iter([_ANY_AIRBYTE_MESSAGE])
        stream_instances_to_read_from = [self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        handler.start_next_partition_generator()

        sentinel = PartitionGenerationCompletedSentinel(self._another_stream)
        messages = list(handler.on_partition_generation_completed(sentinel))

        expected_messages = [
            _ANY_AIRBYTE_MESSAGE,
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    emitted_at=1577836800000.0,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(name=_ANOTHER_STREAM_NAME),
                        status=AirbyteStreamStatus(AirbyteStreamStatus.COMPLETE),
                    ),
                ),
            ),
        ]
        assert messages == expected_messages
        assert in_order_validation_mock.mock_calls.index(
            call._another_stream.cursor.ensure_at_least_one_state_emitted
        ) < in_order_validation_mock.mock_calls.index(call._message_repository.consume_queue)

    def test_handle_partition(self):
        stream_instances_to_read_from = [self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        expected_cursor = handler._stream_name_to_instance[_ANOTHER_STREAM_NAME].cursor

        handler.on_partition(self._a_closed_partition)

        self._thread_pool_manager.submit.assert_called_with(
            self._partition_reader.process_partition, self._a_closed_partition, expected_cursor
        )
        assert (
            self._a_closed_partition in handler._streams_to_running_partitions[_ANOTHER_STREAM_NAME]
        )

    def test_handle_partition_emits_log_message_if_it_should_be_logged(self):
        stream_instances_to_read_from = [self._stream]
        self._slice_logger = Mock(spec=SliceLogger)
        self._slice_logger.should_log_slice_message.return_value = True
        self._slice_logger.create_slice_log_message.return_value = self._log_message

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        expected_cursor = handler._stream_name_to_instance[_STREAM_NAME].cursor

        handler.on_partition(self._an_open_partition)

        self._thread_pool_manager.submit.assert_called_with(
            self._partition_reader.process_partition, self._an_open_partition, expected_cursor
        )
        self._message_repository.emit_message.assert_called_with(self._log_message)

        assert self._an_open_partition in handler._streams_to_running_partitions[_STREAM_NAME]

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_handle_on_partition_complete_sentinel_with_messages_from_repository(self):
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        log_message = Mock(spec=LogMessage)
        partition.to_slice.return_value = log_message
        partition.stream_name.return_value = _STREAM_NAME

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        handler.start_next_partition_generator()
        handler.on_partition(partition)

        sentinel = PartitionCompleteSentinel(partition)

        self._message_repository.consume_queue.return_value = [
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            )
        ]

        messages = list(handler.on_partition_complete_sentinel(sentinel))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            )
        ]
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_handle_on_partition_complete_sentinel_yields_status_message_if_the_stream_is_done(
        self,
    ):
        self._streams_currently_generating_partitions = [self._another_stream]
        stream_instances_to_read_from = [self._another_stream]
        log_message = Mock(spec=LogMessage)
        self._a_closed_partition.to_slice.return_value = log_message
        self._message_repository.consume_queue.return_value = []

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        handler.start_next_partition_generator()
        handler.on_partition(self._a_closed_partition)
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._another_stream)
            )
        )

        sentinel = PartitionCompleteSentinel(self._a_closed_partition)

        messages = list(handler.on_partition_complete_sentinel(sentinel))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(
                            name=_ANOTHER_STREAM_NAME,
                        ),
                        status=AirbyteStreamStatus.COMPLETE,
                    ),
                    emitted_at=1577836800000.0,
                ),
            )
        ]
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_handle_on_partition_complete_sentinel_yields_no_status_message_if_the_stream_is_not_done(
        self,
    ):
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        log_message = Mock(spec=LogMessage)
        partition.to_slice.return_value = log_message
        partition.stream_name.return_value = _STREAM_NAME

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )
        handler.start_next_partition_generator()

        sentinel = PartitionCompleteSentinel(partition)

        messages = list(handler.on_partition_complete_sentinel(sentinel))

        expected_messages = []
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_record_no_status_message_no_repository_messge(self):
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        log_message = Mock(spec=LogMessage)
        partition.to_slice.return_value = log_message
        partition.stream_name.return_value = _STREAM_NAME
        self._message_repository.consume_queue.return_value = []

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Simulate a first record
        list(handler.on_record(self._record))

        messages = list(handler.on_record(self._record))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.RECORD,
                record=AirbyteRecordMessage(
                    stream=_STREAM_NAME,
                    data=self._record_data,
                    emitted_at=1577836800000,
                ),
            )
        ]
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_record_with_repository_messge(self):
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        log_message = Mock(spec=LogMessage)
        partition.to_slice.return_value = log_message
        partition.stream_name.return_value = _STREAM_NAME
        slice_logger = Mock(spec=SliceLogger)
        slice_logger.should_log_slice_message.return_value = True
        slice_logger.create_slice_log_message.return_value = log_message
        self._message_repository.consume_queue.return_value = [
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            )
        ]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        stream = Mock(spec=AbstractStream)
        stream.name = _STREAM_NAME
        stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )

        # Simulate a first record
        list(handler.on_record(self._record))

        messages = list(handler.on_record(self._record))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.RECORD,
                record=AirbyteRecordMessage(
                    stream=_STREAM_NAME,
                    data=self._record_data,
                    emitted_at=1577836800000,
                ),
            ),
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            ),
        ]
        assert messages == expected_messages
        assert handler._record_counter[_STREAM_NAME] == 2

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_record_emits_status_message_on_first_record_no_repository_message(self):
        self._streams_currently_generating_partitions = [_STREAM_NAME]
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        partition.stream_name.return_value = _STREAM_NAME

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        messages = list(handler.on_record(self._record))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    emitted_at=1577836800000.0,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(name=_STREAM_NAME),
                        status=AirbyteStreamStatus(AirbyteStreamStatus.RUNNING),
                    ),
                ),
            ),
            AirbyteMessage(
                type=MessageType.RECORD,
                record=AirbyteRecordMessage(
                    stream=_STREAM_NAME,
                    data=self._record_data,
                    emitted_at=1577836800000,
                ),
            ),
        ]
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_record_emits_status_message_on_first_record_with_repository_message(self):
        stream_instances_to_read_from = [self._stream]
        partition = Mock(spec=Partition)
        log_message = Mock(spec=LogMessage)
        partition.to_slice.return_value = log_message
        partition.stream_name.return_value = _STREAM_NAME
        self._message_repository.consume_queue.return_value = [
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            )
        ]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        stream = Mock(spec=AbstractStream)
        stream.name = _STREAM_NAME
        stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )

        messages = list(handler.on_record(self._record))

        expected_messages = [
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    emitted_at=1577836800000.0,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(name=_STREAM_NAME),
                        status=AirbyteStreamStatus(AirbyteStreamStatus.RUNNING),
                    ),
                ),
            ),
            AirbyteMessage(
                type=MessageType.RECORD,
                record=AirbyteRecordMessage(
                    stream=_STREAM_NAME,
                    data=self._record_data,
                    emitted_at=1577836800000,
                ),
            ),
            AirbyteMessage(
                type=MessageType.LOG,
                log=AirbyteLogMessage(
                    level=LogLevel.INFO, message="message emitted from the repository"
                ),
            ),
        ]
        assert messages == expected_messages

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_exception_return_trace_message_and_on_stream_complete_return_stream_status(self):
        stream_instances_to_read_from = [self._stream, self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._stream)
            )
        )
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._another_stream)
            )
        )

        another_stream = Mock(spec=AbstractStream)
        another_stream.name = _STREAM_NAME
        another_stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_ANOTHER_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )

        exception = StreamThreadException(RuntimeError("Something went wrong"), _STREAM_NAME)

        exception_messages = list(handler.on_exception(exception))
        assert len(exception_messages) == 1
        assert "RuntimeError" in exception_messages[0].trace.error.stack_trace
        assert (
            exception_messages[0].trace.error.message
            == f"An unexpected error occurred in stream {_STREAM_NAME}: RuntimeError"
        )

        assert list(
            handler.on_partition_complete_sentinel(
                PartitionCompleteSentinel(self._an_open_partition)
            )
        ) == [
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    emitted_at=1577836800000.0,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(name=_STREAM_NAME),
                        status=AirbyteStreamStatus(AirbyteStreamStatus.INCOMPLETE),
                    ),
                ),
            )
        ]
        with pytest.raises(AirbyteTracedException):
            handler.is_done()

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_given_underlying_exception_is_traced_exception_on_exception_return_trace_message_and_on_stream_complete_return_stream_status(
        self,
    ):
        stream_instances_to_read_from = [self._stream, self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._stream)
            )
        )
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._another_stream)
            )
        )

        another_stream = Mock(spec=AbstractStream)
        another_stream.name = _STREAM_NAME
        another_stream.as_airbyte_stream.return_value = AirbyteStream(
            name=_ANOTHER_STREAM_NAME,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )

        underlying_exception = AirbyteTracedException()
        exception = StreamThreadException(underlying_exception, _STREAM_NAME)

        exception_messages = list(handler.on_exception(exception))
        assert len(exception_messages) == 1
        assert "AirbyteTracedException" in exception_messages[0].trace.error.stack_trace

        assert list(
            handler.on_partition_complete_sentinel(
                PartitionCompleteSentinel(self._an_open_partition)
            )
        ) == [
            AirbyteMessage(
                type=MessageType.TRACE,
                trace=AirbyteTraceMessage(
                    type=TraceType.STREAM_STATUS,
                    emitted_at=1577836800000.0,
                    stream_status=AirbyteStreamStatusTraceMessage(
                        stream_descriptor=StreamDescriptor(name=_STREAM_NAME),
                        status=AirbyteStreamStatus(AirbyteStreamStatus.INCOMPLETE),
                    ),
                ),
            )
        ]
        with pytest.raises(AirbyteTracedException):
            handler.is_done()

    def test_given_partition_completion_is_not_success_then_do_not_close_partition(self):
        stream_instances_to_read_from = [self._stream, self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._stream)
            )
        )

        list(
            handler.on_partition_complete_sentinel(
                PartitionCompleteSentinel(self._an_open_partition, not _IS_SUCCESSFUL)
            )
        )

        assert self._stream.cursor.close_partition.call_count == 0

    def test_is_done_is_false_if_there_are_any_instances_to_read_from(self):
        stream_instances_to_read_from = [self._stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        assert not handler.is_done()

    def test_is_done_is_false_if_there_are_streams_still_generating_partitions(self):
        stream_instances_to_read_from = [self._stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()

        assert not handler.is_done()

    def test_is_done_is_false_if_all_partitions_are_not_closed(self):
        stream_instances_to_read_from = [self._stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)
        handler.on_partition_generation_completed(
            PartitionGenerationCompletedSentinel(self._stream)
        )

        assert not handler.is_done()

    def test_is_done_is_true_if_all_partitions_are_closed_and_no_streams_are_generating_partitions_and_none_are_still_to_run(
        self,
    ):
        stream_instances_to_read_from = []

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        assert handler.is_done()

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_on_exception_non_ate_uses_templated_message_with_correct_failure_type(self):
        """Regression test: non-ATE exceptions on Path B produce a safe templated message, not the generic fallback."""
        stream_instances_to_read_from = [self._stream, self._another_stream]

        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        handler.start_next_partition_generator()
        handler.on_partition(self._an_open_partition)
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._stream)
            )
        )
        list(
            handler.on_partition_generation_completed(
                PartitionGenerationCompletedSentinel(self._another_stream)
            )
        )

        inner_exception = ValueError("some internal detail: SELECT * FROM secrets")
        exception = StreamThreadException(inner_exception, _STREAM_NAME)

        exception_messages = list(handler.on_exception(exception))
        assert len(exception_messages) == 1

        trace_error = exception_messages[0].trace.error
        # User-facing message uses safe template (no raw exception text)
        assert (
            trace_error.message
            == f"An unexpected error occurred in stream {_STREAM_NAME}: ValueError"
        )
        # Stack trace comes from the inner exception, not the StreamThreadException wrapper
        assert "ValueError" in trace_error.stack_trace
        assert "StreamThreadException" not in trace_error.stack_trace
        # failure_type defaults to system_error for unhandled exceptions
        assert trace_error.failure_type == FailureType.system_error

    @freezegun.freeze_time("2020-01-01T00:00:00")
    def test_start_next_partition_generator(self):
        stream_instances_to_read_from = [self._stream]
        handler = ConcurrentReadProcessor(
            stream_instances_to_read_from,
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        status_message = handler.start_next_partition_generator()

        assert status_message == AirbyteMessage(
            type=MessageType.TRACE,
            trace=AirbyteTraceMessage(
                type=TraceType.STREAM_STATUS,
                emitted_at=1577836800000.0,
                stream_status=AirbyteStreamStatusTraceMessage(
                    stream_descriptor=StreamDescriptor(name=_STREAM_NAME),
                    status=AirbyteStreamStatus(AirbyteStreamStatus.STARTED),
                ),
            ),
        )

        assert _STREAM_NAME in handler._streams_currently_generating_partitions
        self._thread_pool_manager.submit.assert_called_with(
            self._partition_enqueuer.generate_partitions, self._stream
        )


class TestBlockSimultaneousRead(unittest.TestCase):
    """Tests for block_simultaneous_read functionality"""

    def setUp(self):
        self._partition_enqueuer = Mock(spec=PartitionEnqueuer)
        self._thread_pool_manager = Mock(spec=ThreadPoolManager)
        self._logger = Mock(spec=logging.Logger)
        self._slice_logger = Mock(spec=SliceLogger)
        self._message_repository = Mock(spec=MessageRepository)
        self._message_repository.consume_queue.return_value = []
        self._partition_reader = Mock(spec=PartitionReader)

    def _create_mock_stream(self, name: str, block_simultaneous_read: str = ""):
        """Helper to create a mock stream"""
        stream = Mock(spec=AbstractStream)
        stream.name = name
        stream.block_simultaneous_read = block_simultaneous_read
        stream.as_airbyte_stream.return_value = AirbyteStream(
            name=name,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )
        stream.cursor.ensure_at_least_one_state_emitted = Mock()
        return stream

    def _create_mock_stream_with_parent(
        self, name: str, parent_stream, block_simultaneous_read: str = ""
    ):
        """Helper to create a mock stream with a parent stream."""
        stream = Mock(spec=DefaultStream)
        stream.name = name
        stream.block_simultaneous_read = block_simultaneous_read
        stream.as_airbyte_stream.return_value = AirbyteStream(
            name=name,
            json_schema={},
            supported_sync_modes=[SyncMode.full_refresh],
        )
        stream.cursor.ensure_at_least_one_state_emitted = Mock()

        mock_partition_router = Mock(spec=SubstreamPartitionRouter)
        mock_parent_config = Mock()
        mock_parent_config.stream = parent_stream
        mock_partition_router.parent_stream_configs = [mock_parent_config]
        stream.get_partition_router.return_value = mock_partition_router

        return stream

    def test_defer_stream_when_self_active(self):
        """Test that a stream is deferred when it's already active"""
        stream = self._create_mock_stream("stream1", block_simultaneous_read="api_group")

        handler = ConcurrentReadProcessor(
            [stream],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark stream as active
        handler._active_stream_names.add("stream1")

        # Try to start the stream again
        result = handler.start_next_partition_generator()

        # Should return None (no stream started)
        assert result is None

        # Stream should be back in the queue
        assert len(handler._stream_instances_to_start_partition_generation) == 1
        assert handler._stream_instances_to_start_partition_generation[0] == stream

        # Logger should have been called to log deferral
        assert any(
            "Deferring stream 'stream1' (group 'api_group') because it's already active"
            in str(call)
            for call in self._logger.info.call_args_list
        )

    def test_defer_stream_when_parent_active(self):
        """Test that a stream is deferred when its parent is active"""
        parent_stream = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child_stream = self._create_mock_stream_with_parent(
            "child", parent_stream, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent_stream, child_stream],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark parent as active
        handler._active_stream_names.add("parent")

        # Remove parent from queue (simulate it's already started)
        handler._stream_instances_to_start_partition_generation = [child_stream]

        # Try to start child
        result = handler.start_next_partition_generator()

        # Should return None (child deferred)
        assert result is None

        # Child should be back in the queue
        assert len(handler._stream_instances_to_start_partition_generation) == 1
        assert handler._stream_instances_to_start_partition_generation[0] == child_stream

        # Logger should have been called
        assert any(
            "Deferring stream 'child' because parent stream(s)" in str(call)
            for call in self._logger.info.call_args_list
        )

    def test_defer_stream_when_grandparent_active(self):
        """Test that a stream is deferred when its grandparent is active"""
        grandparent = self._create_mock_stream("grandparent", block_simultaneous_read="api_group")
        parent = self._create_mock_stream_with_parent(
            "parent", grandparent, block_simultaneous_read="api_group"
        )
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [grandparent, parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark grandparent as active
        handler._active_stream_names.add("grandparent")

        # Only child in queue
        handler._stream_instances_to_start_partition_generation = [child]

        # Try to start child
        result = handler.start_next_partition_generator()

        # Should return None (child deferred because grandparent is active)
        assert result is None

        # Child should be back in the queue
        assert len(handler._stream_instances_to_start_partition_generation) == 1

    def test_different_groups_do_not_block_each_other(self):
        """Test that independent streams with different groups don't block each other"""
        stream1 = self._create_mock_stream("stream1", block_simultaneous_read="group1")
        stream2 = self._create_mock_stream("stream2", block_simultaneous_read="group2")

        handler = ConcurrentReadProcessor(
            [stream1, stream2],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Start stream1
        handler.start_next_partition_generator()
        assert "stream1" in handler._active_stream_names

        # Stream2 should start successfully even though stream1 is active
        # because they're in different groups
        result = handler.start_next_partition_generator()

        # Should start stream2 (different group, no blocking)
        assert result is not None
        assert "stream2" in handler._active_stream_names
        assert len(handler._stream_instances_to_start_partition_generation) == 0

    def test_retry_blocked_stream_after_partition_generation(self):
        """Test that blocked stream is retried after partition generation completes"""
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Start parent
        handler.start_next_partition_generator()
        assert "parent" in handler._active_stream_names

        # Complete partition generation for parent (parent has no partitions, so it's done)
        sentinel = PartitionGenerationCompletedSentinel(parent)
        messages = list(handler.on_partition_generation_completed(sentinel))

        # Child should have been started automatically by on_partition_generation_completed
        # (it calls start_next_partition_generator internally)
        assert "child" in handler._active_stream_names

        # Parent should be RE-ACTIVATED because child needs to read from it during partition generation
        # This is the correct behavior - prevents simultaneous reads of parent
        assert "parent" in handler._active_stream_names

        # Verify the queue is now empty (both streams were started)
        assert len(handler._stream_instances_to_start_partition_generation) == 0

    def test_blocked_stream_added_to_end_of_queue(self):
        """Test that blocked streams are added to the end of the queue"""
        stream1 = self._create_mock_stream("stream1", block_simultaneous_read="api_group")
        stream2 = self._create_mock_stream("stream2", block_simultaneous_read="")
        stream3 = self._create_mock_stream("stream3", block_simultaneous_read="")

        handler = ConcurrentReadProcessor(
            [stream1, stream2, stream3],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark stream1 as active
        handler._active_stream_names.add("stream1")

        # Try to start streams in order: stream1, stream2, stream3
        result1 = handler.start_next_partition_generator()

        # stream1 should be deferred, stream2 should start
        assert result1 is not None
        assert "stream2" in handler._active_stream_names

        # Queue should now be [stream3, stream1] (stream1 moved to end)
        assert len(handler._stream_instances_to_start_partition_generation) == 2
        assert handler._stream_instances_to_start_partition_generation[0] == stream3
        assert handler._stream_instances_to_start_partition_generation[1] == stream1

    def test_no_defer_when_flag_false(self):
        """Test that blocking doesn't occur when block_simultaneous_read="" """
        stream = self._create_mock_stream("stream1", block_simultaneous_read="")

        handler = ConcurrentReadProcessor(
            [stream],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark stream as active
        handler._active_stream_names.add("stream1")

        # Try to start the stream again (should succeed because flag is False)
        result = handler.start_next_partition_generator()

        # Should return a status message (stream started)
        assert result is not None
        assert isinstance(result, AirbyteMessage)

        # Queue should be empty
        assert len(handler._stream_instances_to_start_partition_generation) == 0

    def test_collect_parent_streams_multi_level(self):
        """Test that _collect_all_parent_stream_names works recursively"""
        grandparent = self._create_mock_stream("grandparent")
        parent = self._create_mock_stream_with_parent("parent", grandparent)
        child = self._create_mock_stream_with_parent("child", parent)

        handler = ConcurrentReadProcessor(
            [grandparent, parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Collect parents for child
        parents = handler._collect_all_parent_stream_names("child")

        # Should include both parent and grandparent
        assert "parent" in parents
        assert "grandparent" in parents
        assert len(parents) == 2

    def test_deactivate_parents_when_partition_generation_completes(self):
        """Test that parent streams are deactivated when partition generation completes"""
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Manually mark both as active (simulating partition generation for child)
        handler._active_stream_names.add("parent")
        handler._active_stream_names.add("child")
        handler._streams_currently_generating_partitions.append("child")

        # Ensure child has running partitions (so it doesn't trigger _on_stream_is_done)
        mock_partition = Mock(spec=Partition)
        mock_partition.stream_name.return_value = "child"
        handler._streams_to_running_partitions["child"] = {mock_partition}

        # Remove both streams from the queue so start_next_partition_generator doesn't start them
        # This simulates the scenario where both streams have already been started
        handler._stream_instances_to_start_partition_generation = []

        # Complete partition generation for child
        sentinel = PartitionGenerationCompletedSentinel(child)

        list(handler.on_partition_generation_completed(sentinel))

        # Parent should be deactivated (it was only needed for partition generation)
        assert "parent" not in handler._active_stream_names

        # Child should still be active (it's reading records)
        assert "child" in handler._active_stream_names

    def test_deactivate_only_stream_when_done(self):
        """Test that only the stream itself is deactivated when done, not parents"""
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark both as active
        handler._active_stream_names.add("parent")
        handler._active_stream_names.add("child")

        # Start child and mark it as done
        handler._stream_instances_to_start_partition_generation = []
        handler._streams_currently_generating_partitions = []
        handler._streams_to_running_partitions["child"] = set()

        # Call _on_stream_is_done for child
        list(handler._on_stream_is_done("child"))

        # Child should be deactivated
        assert "child" not in handler._active_stream_names

        # Parent should still be active (not deactivated)
        assert "parent" in handler._active_stream_names

    def test_multiple_blocked_streams_retry_in_order(self):
        """Test that multiple blocked streams are retried in order"""
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child1 = self._create_mock_stream_with_parent(
            "child1", parent, block_simultaneous_read="api_group"
        )
        child2 = self._create_mock_stream_with_parent(
            "child2", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child1, child2],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Start parent
        result = handler.start_next_partition_generator()
        assert result is not None
        assert "parent" in handler._active_stream_names
        assert "api_group" in handler._active_groups
        assert "parent" in handler._active_groups["api_group"]

        # Try to start next stream (child1) - should be deferred because parent is active
        result = handler.start_next_partition_generator()
        assert result is None  # child1 was deferred

        # After first deferral, we should still have 2 streams in queue (child1 moved to end)
        assert len(handler._stream_instances_to_start_partition_generation) == 2
        # child1 was moved to the back, so the queue has the other child first
        queue_streams = handler._stream_instances_to_start_partition_generation
        assert child1 in queue_streams
        assert child2 in queue_streams

        # Try to start next stream (child2) - should also be deferred
        result = handler.start_next_partition_generator()
        assert result is None  # child2 was deferred

        # Both streams still in queue, but order may have changed
        assert len(handler._stream_instances_to_start_partition_generation) == 2

        # Verify neither child is active yet (both blocked by parent)
        assert "child1" not in handler._active_stream_names
        assert "child2" not in handler._active_stream_names

        # Verify deferral was logged for both children
        logger_calls = [str(call) for call in self._logger.info.call_args_list]
        assert any("Deferring stream 'child1'" in call for call in logger_calls)
        assert any("Deferring stream 'child2'" in call for call in logger_calls)

        # Simulate parent completing partition generation (parent has no partitions, so it's done)
        sentinel = PartitionGenerationCompletedSentinel(parent)
        list(handler.on_partition_generation_completed(sentinel))

        # After parent completes, one of the children should start (whichever was first in queue)
        # We know at least one child started because the queue shrunk
        assert len(handler._stream_instances_to_start_partition_generation) == 1

        # Verify that exactly one child is now active
        children_active = [
            name for name in ["child1", "child2"] if name in handler._active_stream_names
        ]
        assert len(children_active) == 1, (
            f"Expected exactly one child active, got: {children_active}"
        )

        # Parent should be re-activated because the active child needs to read from it
        assert "parent" in handler._active_stream_names

    def test_child_without_flag_blocked_by_parent_with_flag(self):
        """Test that a child WITHOUT block_simultaneous_read is blocked by parent WITH the flag"""
        # Parent has the flag, child does NOT
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child = self._create_mock_stream_with_parent("child", parent, block_simultaneous_read="")

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark parent as active and already started (remove from queue)
        handler._active_stream_names.add("parent")
        handler._stream_instances_to_start_partition_generation.remove(parent)

        # Try to start child (should be deferred even though child doesn't have the flag)
        result = handler.start_next_partition_generator()

        # Child should be deferred because parent has block_simultaneous_read="api_group" and is active
        assert result is None  # No stream started
        assert "child" not in handler._active_stream_names
        # Child should be moved to end of queue (still 1 stream in queue)
        assert len(handler._stream_instances_to_start_partition_generation) == 1
        assert handler._stream_instances_to_start_partition_generation[0] == child

    def test_child_with_flag_not_blocked_by_parent_without_flag(self):
        """Test that a child WITH block_simultaneous_read is NOT blocked by parent WITHOUT the flag"""
        # Parent does NOT have the flag, child does
        parent = self._create_mock_stream("parent", block_simultaneous_read="")
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Mark parent as active and already started (remove from queue)
        handler._active_stream_names.add("parent")
        handler._stream_instances_to_start_partition_generation.remove(parent)

        # Try to start child (should succeed even though parent is active)
        result = handler.start_next_partition_generator()

        # Child should start successfully because parent doesn't have block_simultaneous_read
        assert result is not None  # Stream started
        assert "child" in handler._active_stream_names
        # Queue should now be empty (both streams started)
        assert len(handler._stream_instances_to_start_partition_generation) == 0

    def test_unrelated_streams_in_same_group_block_each_other(self):
        """Test that multiple unrelated streams with the same group name block each other"""
        # Create three unrelated streams (no parent-child relationship) in the same group
        stream1 = self._create_mock_stream("stream1", block_simultaneous_read="shared_endpoint")
        stream2 = self._create_mock_stream("stream2", block_simultaneous_read="shared_endpoint")
        stream3 = self._create_mock_stream("stream3", block_simultaneous_read="shared_endpoint")

        handler = ConcurrentReadProcessor(
            [stream1, stream2, stream3],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Start stream1
        result = handler.start_next_partition_generator()
        assert result is not None
        assert "stream1" in handler._active_stream_names
        assert "shared_endpoint" in handler._active_groups
        assert "stream1" in handler._active_groups["shared_endpoint"]

        # Try to start stream2 (should be deferred because it's in the same group)
        result = handler.start_next_partition_generator()
        # stream2 should be deferred, stream3 should also be deferred
        # All three are in same group, only stream1 is active
        assert result is None  # No stream started

        # Both stream2 and stream3 should be in the queue
        assert len(handler._stream_instances_to_start_partition_generation) == 2

        # Verify logger was called with deferral message
        assert any(
            "Deferring stream 'stream2'" in str(call) and "shared_endpoint" in str(call)
            for call in self._logger.info.call_args_list
        )

    def test_child_starts_after_parent_completes_via_partition_complete_sentinel(self):
        """Test that child stream starts after parent completes via on_partition_complete_sentinel"""
        parent = self._create_mock_stream("parent", block_simultaneous_read="api_group")
        child = self._create_mock_stream_with_parent(
            "child", parent, block_simultaneous_read="api_group"
        )

        handler = ConcurrentReadProcessor(
            [parent, child],
            self._partition_enqueuer,
            self._thread_pool_manager,
            self._logger,
            self._slice_logger,
            self._message_repository,
            self._partition_reader,
        )

        # Start parent
        handler.start_next_partition_generator()
        assert "parent" in handler._active_stream_names

        # Try to start child (should be deferred)
        result = handler.start_next_partition_generator()
        assert result is None
        assert "child" not in handler._active_stream_names
        assert len(handler._stream_instances_to_start_partition_generation) == 1

        # Create a partition for parent and add it to running partitions
        # (parent is already in _streams_currently_generating_partitions from start_next_partition_generator)
        mock_partition = Mock(spec=Partition)
        mock_partition.stream_name.return_value = "parent"
        handler._streams_to_running_partitions["parent"].add(mock_partition)

        # Complete partition generation for parent
        sentinel_gen = PartitionGenerationCompletedSentinel(parent)
        list(handler.on_partition_generation_completed(sentinel_gen))

        # Now complete the partition (this triggers stream done)
        sentinel_complete = PartitionCompleteSentinel(mock_partition)
        messages = list(handler.on_partition_complete_sentinel(sentinel_complete))

        # Child should have been started automatically
        assert "child" in handler._active_stream_names
        assert len(handler._stream_instances_to_start_partition_generation) == 0

        # Verify a STARTED message was emitted for child
        started_messages = [
            msg
            for msg in messages
            if msg.type == MessageType.TRACE
            and msg.trace.stream_status
            and msg.trace.stream_status.status == AirbyteStreamStatus.STARTED
        ]
        assert len(started_messages) == 1
        assert started_messages[0].trace.stream_status.stream_descriptor.name == "child"


def test_is_done_raises_when_partition_generation_queue_not_empty():
    """Test is_done raises AirbyteTracedException if streams remain in the partition generation queue."""
    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 = ""
    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,
    )

    # Artificially mark the stream as done without removing it from the partition generation queue
    handler._streams_done.add("stuck_stream")

    with pytest.raises(
        AirbyteTracedException,
        match="Partition generation queue is not empty after all streams completed",
    ):
        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="Active stream groups are not empty after all streams completed",
    ):
        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"}
