teng-lin--notebooklm-py/tests/unit/test_source_status.py

"""Unit tests for source status and polling functionality."""

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

import notebooklm._sources as _sources_mod
from notebooklm._sources import SourcesAPI
from notebooklm.types import (
    Source,
    SourceError,
    SourceNotFoundError,
    SourceProcessingError,
    SourceStatus,
    SourceTimeoutError,
)


class TestSourceStatus:
    """Tests for SourceStatus enum."""

    def test_status_values(self):
        """Test SourceStatus enum values."""
        assert SourceStatus.PROCESSING == 1
        assert SourceStatus.READY == 2
        assert SourceStatus.ERROR == 3

    def test_status_is_int_enum(self):
        """Test that SourceStatus values can be compared with ints."""
        assert SourceStatus.READY == 2
        assert SourceStatus.READY == 2


class TestSourceStatusProperties:
    """Tests for Source status properties."""

    def test_is_ready_when_ready(self):
        """Test is_ready returns True for READY status."""
        source = Source(id="src_1", status=SourceStatus.READY)
        assert source.is_ready is True
        assert source.is_processing is False
        assert source.is_error is False

    def test_is_processing_when_processing(self):
        """Test is_processing returns True for PROCESSING status."""
        source = Source(id="src_1", status=SourceStatus.PROCESSING)
        assert source.is_ready is False
        assert source.is_processing is True
        assert source.is_error is False

    def test_is_error_when_error(self):
        """Test is_error returns True for ERROR status."""
        source = Source(id="src_1", status=SourceStatus.ERROR)
        assert source.is_ready is False
        assert source.is_processing is False
        assert source.is_error is True

    def test_default_status_is_ready(self):
        """Test that default status is READY."""
        source = Source(id="src_1")
        assert source.status == SourceStatus.READY
        assert source.is_ready is True


class TestSourceExceptions:
    """Tests for source exception classes."""

    def test_source_error_is_base(self):
        """Test SourceError is the base exception."""
        assert issubclass(SourceProcessingError, SourceError)
        assert issubclass(SourceTimeoutError, SourceError)
        assert issubclass(SourceNotFoundError, SourceError)

    def test_source_processing_error(self):
        """Test SourceProcessingError attributes."""
        error = SourceProcessingError("src_123", status=3)
        assert error.source_id == "src_123"
        assert error.status == 3
        assert "src_123" in str(error)

    def test_source_timeout_error(self):
        """Test SourceTimeoutError attributes."""
        error = SourceTimeoutError("src_123", timeout=60.0, last_status=1)
        assert error.source_id == "src_123"
        assert error.timeout == 60.0
        assert error.last_status == 1
        assert "60.0" in str(error)
        assert "src_123" in str(error)

    def test_source_not_found_error(self):
        """Test SourceNotFoundError attributes."""
        error = SourceNotFoundError("src_123")
        assert error.source_id == "src_123"
        assert "src_123" in str(error)


class TestWaitUntilReady:
    """Tests for wait_until_ready method."""

    @pytest.fixture
    def sources_api(self):
        """Create a SourcesAPI with mocked core."""
        core = MagicMock()
        return SourcesAPI(core, uploader=MagicMock())

    @pytest.mark.asyncio
    async def test_returns_immediately_if_ready(self, sources_api):
        """Test that wait_until_ready returns immediately if source is ready."""
        ready_source = Source(id="src_1", title="Test", status=SourceStatus.READY)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = ready_source

            result = await sources_api.wait_until_ready("nb_1", "src_1", timeout=10.0)

            assert result.is_ready
            assert mock_get.call_count == 1

    @pytest.mark.asyncio
    async def test_polls_until_ready(self, sources_api):
        """Test that wait_until_ready polls until source becomes ready."""
        processing_source = Source(id="src_1", status=SourceStatus.PROCESSING)
        ready_source = Source(id="src_1", status=SourceStatus.READY)

        call_count = 0

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count < 3:
                return processing_source
            return ready_source

        with patch.object(sources_api, "get_or_none", side_effect=mock_get):
            result = await sources_api.wait_until_ready(
                "nb_1", "src_1", timeout=10.0, initial_interval=0.01
            )

            assert result.is_ready
            assert call_count == 3

    @pytest.mark.asyncio
    async def test_raises_processing_error_on_error_status(self, sources_api):
        """Test that wait_until_ready raises SourceProcessingError on ERROR status."""
        # Use a non-audio terminal type (e.g. PDF, type_code=3) so the new
        # audio-tolerance gate doesn't keep polling and trigger a timeout.
        error_source = Source(id="src_1", status=SourceStatus.ERROR, _type_code=3)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = error_source

            with pytest.raises(SourceProcessingError) as exc_info:
                await sources_api.wait_until_ready("nb_1", "src_1", timeout=10.0)

            assert exc_info.value.source_id == "src_1"
            assert exc_info.value.status == SourceStatus.ERROR

    @pytest.mark.asyncio
    async def test_wait_until_ready_tolerates_transient_status3_for_audio(self, sources_api):
        """Audio sources (type_code=10) briefly report status=3 during transcription
        before settling at status=2. wait_until_ready should keep polling instead of
        raising SourceProcessingError on the first status=3 observation.
        """
        transient_error = Source(id="src_audio", status=SourceStatus.ERROR, _type_code=10)
        ready = Source(id="src_audio", status=SourceStatus.READY, _type_code=10)

        call_count = 0

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return transient_error
            return ready

        with patch.object(sources_api, "get_or_none", side_effect=mock_get):
            result = await sources_api.wait_until_ready(
                "nb_1", "src_audio", timeout=10.0, initial_interval=0.01
            )

        assert result.is_ready
        assert call_count == 2

    @pytest.mark.asyncio
    @pytest.mark.parametrize("type_code", [None, 0], ids=["none", "zero"])
    async def test_wait_until_ready_tolerates_transient_status3_for_unclassified(
        self, sources_api, type_code
    ):
        """Sources with no type code yet (None / 0) can also briefly report
        status=3 during processing. Keep polling rather than raising.
        """
        transient_error = Source(id="src_x", status=SourceStatus.ERROR, _type_code=type_code)
        # Once the source is classified as audio (type 10) and ready,
        # the loop should return it.
        ready = Source(id="src_x", status=SourceStatus.READY, _type_code=10)

        call_count = 0

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return transient_error
            return ready

        with patch.object(sources_api, "get_or_none", side_effect=mock_get):
            result = await sources_api.wait_until_ready(
                "nb_1", "src_x", timeout=10.0, initial_interval=0.01
            )

        assert result.is_ready, f"expected ready for type_code={type_code!r}"

    @pytest.mark.asyncio
    async def test_wait_until_ready_still_raises_on_pdf_status3(self, sources_api):
        """Non-audio terminal types (e.g. PDF type_code=3) must still raise
        SourceProcessingError on status=3 — the audio tolerance is gated by
        type_code and should not regress the existing error-detection path.
        """
        pdf_error = Source(id="src_pdf", status=SourceStatus.ERROR, _type_code=3)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = pdf_error

            with pytest.raises(SourceProcessingError) as exc_info:
                await sources_api.wait_until_ready("nb_1", "src_pdf", timeout=10.0)

            assert exc_info.value.source_id == "src_pdf"
            assert exc_info.value.status == SourceStatus.ERROR

    @pytest.mark.asyncio
    async def test_raises_not_found_error_when_source_missing(self, sources_api):
        """Test that wait_until_ready raises SourceNotFoundError when source not found."""
        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = None

            with pytest.raises(SourceNotFoundError) as exc_info:
                await sources_api.wait_until_ready("nb_1", "src_1", timeout=10.0)

            assert exc_info.value.source_id == "src_1"

    @pytest.mark.asyncio
    async def test_raises_timeout_error(self, sources_api):
        """Test that wait_until_ready raises SourceTimeoutError on timeout."""
        processing_source = Source(id="src_1", status=SourceStatus.PROCESSING)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = processing_source

            with pytest.raises(SourceTimeoutError) as exc_info:
                await sources_api.wait_until_ready(
                    "nb_1", "src_1", timeout=0.05, initial_interval=0.02
                )

            assert exc_info.value.source_id == "src_1"
            assert exc_info.value.last_status == SourceStatus.PROCESSING

    @pytest.mark.asyncio
    async def test_exponential_backoff(self, sources_api):
        """Test that polling uses exponential backoff."""
        processing_source = Source(id="src_1", status=SourceStatus.PROCESSING)
        ready_source = Source(id="src_1", status=SourceStatus.READY)

        call_count = 0
        sleep_intervals = []

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count < 4:
                return processing_source
            return ready_source

        original_sleep = asyncio.sleep

        async def mock_sleep(delay):
            sleep_intervals.append(delay)
            await original_sleep(0.001)  # Minimal actual sleep

        with (
            patch.object(sources_api, "get_or_none", side_effect=mock_get),
            patch.object(_sources_mod.asyncio, "sleep", side_effect=mock_sleep),
        ):
            await sources_api.wait_until_ready(
                "nb_1",
                "src_1",
                timeout=10.0,
                initial_interval=0.05,
                backoff_factor=2.0,
                max_interval=1.0,
            )

        # Check that intervals increase exponentially
        assert len(sleep_intervals) >= 2
        # With initial_interval=0.05 and backoff_factor=2.0:
        # First sleep: 0.05, Second sleep: 0.10, Third sleep: 0.20
        assert sleep_intervals[1] >= sleep_intervals[0] * 1.5


class TestWaitUntilRegistered:
    """Tests for wait_until_registered helper (narrow forced-rename wait)."""

    @pytest.fixture
    def sources_api(self):
        core = MagicMock()
        return SourcesAPI(core, uploader=MagicMock())

    @pytest.mark.asyncio
    async def test_wait_until_registered_returns_on_processing(self, sources_api):
        """Status=PROCESSING (1) on first poll → return immediately."""
        processing = Source(id="src_1", status=SourceStatus.PROCESSING, _type_code=8)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = processing

            result = await sources_api.wait_until_registered("nb_1", "src_1", timeout=10.0)

            assert result is processing
            assert mock_get.call_count == 1

    @pytest.mark.asyncio
    async def test_wait_until_registered_returns_on_ready(self, sources_api):
        """Status=READY (2) on first poll → return immediately."""
        ready = Source(id="src_1", status=SourceStatus.READY, _type_code=8)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = ready

            result = await sources_api.wait_until_registered("nb_1", "src_1", timeout=10.0)

            assert result is ready
            assert mock_get.call_count == 1

    @pytest.mark.asyncio
    async def test_wait_until_registered_polls_when_not_yet_visible(self, sources_api):
        """Source not yet in list → keep polling; return once it appears with status=PROCESSING."""
        processing = Source(id="src_1", status=SourceStatus.PROCESSING, _type_code=8)

        call_count = 0

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return None  # Not yet visible in listing
            return processing

        with patch.object(sources_api, "get_or_none", side_effect=mock_get):
            result = await sources_api.wait_until_registered(
                "nb_1", "src_1", timeout=10.0, initial_interval=0.01
            )

        assert result is processing
        assert call_count == 2

    @pytest.mark.asyncio
    async def test_wait_until_registered_respects_audio_transient(self, sources_api):
        """Audio (type 10) transient ERROR → keep polling; return on next non-error status."""
        transient_error = Source(id="src_audio", status=SourceStatus.ERROR, _type_code=10)
        processing = Source(id="src_audio", status=SourceStatus.PROCESSING, _type_code=10)

        call_count = 0

        async def mock_get(notebook_id, source_id):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return transient_error
            return processing

        with patch.object(sources_api, "get_or_none", side_effect=mock_get):
            result = await sources_api.wait_until_registered(
                "nb_1", "src_audio", timeout=10.0, initial_interval=0.01
            )

        assert result is processing
        assert call_count == 2

    @pytest.mark.asyncio
    async def test_wait_until_registered_raises_on_non_transient_error(self, sources_api):
        """Status=ERROR for a non-transient type (e.g. PDF type_code=3) → raise immediately."""
        pdf_error = Source(id="src_pdf", status=SourceStatus.ERROR, _type_code=3)

        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = pdf_error

            with pytest.raises(SourceProcessingError):
                await sources_api.wait_until_registered("nb_1", "src_pdf", timeout=10.0)

    @pytest.mark.asyncio
    async def test_wait_until_registered_times_out(self, sources_api):
        """If the source never registers, wait_until_registered raises SourceTimeoutError."""
        with patch.object(sources_api, "get_or_none", new_callable=AsyncMock) as mock_get:
            mock_get.return_value = None

            with pytest.raises(SourceTimeoutError):
                await sources_api.wait_until_registered(
                    "nb_1", "src_1", timeout=0.05, initial_interval=0.02
                )


class TestWaitForSources:
    """Tests for wait_for_sources method."""

    @pytest.fixture
    def sources_api(self):
        """Create a SourcesAPI with mocked core."""
        core = MagicMock()
        return SourcesAPI(core, uploader=MagicMock())

    @pytest.mark.asyncio
    async def test_waits_for_multiple_sources(self, sources_api):
        """Test wait_for_sources waits for all sources in parallel."""
        ready_sources = [
            Source(id="src_1", status=SourceStatus.READY),
            Source(id="src_2", status=SourceStatus.READY),
        ]

        async def mock_wait(notebook_id, source_id, **kwargs):
            for s in ready_sources:
                if s.id == source_id:
                    return s
            raise SourceNotFoundError(source_id)

        with patch.object(sources_api, "wait_until_ready", side_effect=mock_wait):
            results = await sources_api.wait_for_sources("nb_1", ["src_1", "src_2"], timeout=10.0)

            assert len(results) == 2
            assert all(s.is_ready for s in results)

    @pytest.mark.asyncio
    async def test_raises_on_any_failure(self, sources_api):
        """Test wait_for_sources raises if any source fails."""

        async def mock_wait(notebook_id, source_id, **kwargs):
            if source_id == "src_2":
                raise SourceProcessingError(source_id, status=3)
            return Source(id=source_id, status=SourceStatus.READY)

        with (
            patch.object(sources_api, "wait_until_ready", side_effect=mock_wait),
            pytest.raises(SourceProcessingError),
        ):
            await sources_api.wait_for_sources("nb_1", ["src_1", "src_2"], timeout=10.0)