/tests/messaging/test_limiter.py

import asyncio
import contextlib
import time
from collections.abc import Callable

import pytest
import pytest_asyncio

from free_claude_code.messaging.limiter import MessagingRateLimiter


class TestMessagingRateLimiter:
    """Tests for MessagingRateLimiter."""

    @pytest_asyncio.fixture(autouse=True)
    async def limiter_factory(self):
        """Build started limiters and stop every instance after the test."""
        instances: list[MessagingRateLimiter] = []

        def create(
            *, rate_limit: int = 1, rate_window: float = 1.0
        ) -> MessagingRateLimiter:
            limiter = MessagingRateLimiter(
                rate_limit=rate_limit,
                rate_window=rate_window,
            )
            limiter.start()
            instances.append(limiter)
            return limiter

        self.create_limiter: Callable[..., MessagingRateLimiter] = create
        yield
        for limiter in reversed(instances):
            await limiter.shutdown(timeout=0.1)

    @pytest.mark.asyncio
    async def test_instances_are_independent(self):
        """Each messaging runtime receives independent limiter state."""
        limiter1 = self.create_limiter(rate_limit=1, rate_window=0.5)
        limiter2 = self.create_limiter(rate_limit=99, rate_window=99.0)

        assert limiter1 is not limiter2
        assert limiter1.limiter._rate_limit == 1
        assert limiter1.limiter._rate_window == 0.5
        assert limiter2.limiter._rate_limit == 99
        assert limiter2.limiter._rate_window == 99.0

        await limiter1.shutdown(timeout=0.1)

        async def succeed() -> str:
            return "still running"

        assert await limiter2.enqueue(succeed) == "still running"

    @pytest.mark.asyncio
    async def test_start_is_required_and_shutdown_is_idempotent(self):
        limiter = MessagingRateLimiter(rate_limit=1, rate_window=1.0)

        async def succeed() -> str:
            return "ok"

        with pytest.raises(RuntimeError, match="has not been started"):
            await limiter.enqueue(succeed)

        limiter.start()
        limiter.start()
        assert await limiter.enqueue(succeed) == "ok"
        await limiter.shutdown(timeout=0.1)
        await limiter.shutdown(timeout=0.1)

        with pytest.raises(RuntimeError, match="is closed"):
            await limiter.enqueue(succeed)

    @pytest.mark.asyncio
    async def test_compaction(self):
        """
        Verify multiple rapid requests with same dedup_key are compacted.
        Logic ported from verify_limiter.py
        """
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        call_counts = {}

        async def mock_edit(msg_id, content):
            call_counts[msg_id] = call_counts.get(msg_id, 0) + 1
            return f"done_{content}"

        # Spam 5 edits
        for i in range(5):
            limiter.fire_and_forget(
                lambda i=i: mock_edit("msg1", f"update_{i}"), dedup_key="edit:msg1"
            )

        # Wait for processing
        # 1st might go through immediately, subsequent ones queue and compact
        await asyncio.sleep(2.5)

        # Expected: ~2 calls (first and last)
        assert call_counts["msg1"] <= 2, (
            f"Expected compaction to reduce calls, but got {call_counts.get('msg1', 0)}"
        )
        assert call_counts["msg1"] >= 1, "Expected at least one call"

    @pytest.mark.asyncio
    async def test_compaction_and_futures_resolution(self):
        """
        Verify that even when compacted, all futures resolve to the result of the LAST execution.
        Logic ported from verify_limiter_v2.py
        """
        limiter = self.create_limiter(rate_limit=1, rate_window=0.5)

        call_counts = {}
        msg_id = "test_msg_hang"

        async def mock_edit(mid, content):
            call_counts[mid] = call_counts.get(mid, 0) + 1
            await asyncio.sleep(0.05)
            return f"result_{content}"

        async def task(i):
            return await limiter.enqueue(
                lambda i=i: mock_edit(msg_id, f"v{i}"), dedup_key=f"edit:{msg_id}"
            )

        start_time = time.time()

        # Enqueue 3 tasks concurrently
        results = await asyncio.gather(task(1), task(2), task(3))

        duration = time.time() - start_time

        # All results should be the LAST one executed
        for res in results:
            assert res == "result_v3", f"Expected result_v3, got {res}"

        # Should be reasonably fast
        assert duration < 2.0, "Execution took too long"

        # Calls should be compacted
        assert call_counts[msg_id] <= 2, f"Too many actual calls: {call_counts[msg_id]}"

    @pytest.mark.asyncio
    async def test_flood_wait_handling(self):
        """Test that FloodWait exceptions pause the worker."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        # Mock exception with .seconds attribute
        class FloodWait(Exception):
            def __init__(self, seconds):
                self.seconds = seconds
                super().__init__(f"Flood wait {seconds}s")

        call_count = 0

        async def mock_fail():
            nonlocal call_count
            call_count += 1
            raise FloodWait(1)  # 1 second wait

        async def mock_success():
            nonlocal call_count
            call_count += 1
            return "success"

        # First call fails and triggers pause
        with contextlib.suppress(Exception):
            await limiter.enqueue(mock_fail, dedup_key="key1")

        assert limiter._paused_until > 0

        # Enqueue success, it should wait
        start = time.time()
        await limiter.enqueue(mock_success, dedup_key="key2")
        duration = time.time() - start

        # Should have waited at least ~1s
        assert duration >= 0.9, (
            f"Should have waited for FloodWait, but took {duration:.2f}s"
        )
        assert call_count == 2

    @pytest.mark.asyncio
    async def test_flood_wait_retry_after_parsing(self):
        """Error message with 'retry after N' parses the wait seconds."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        async def mock_flood():
            raise Exception("Flood wait: retry after 2 seconds")

        with contextlib.suppress(Exception):
            await limiter.enqueue(mock_flood, dedup_key="retry_parse")

        # Should have parsed "after 2" -> 2 seconds
        assert limiter._paused_until > 0

    @pytest.mark.asyncio
    async def test_non_flood_exception_no_pause(self):
        """Non-flood exception doesn't trigger pause."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        async def mock_error():
            raise ValueError("some regular error")

        with contextlib.suppress(ValueError):
            await limiter.enqueue(mock_error, dedup_key="non_flood")

        # Should NOT have paused since it's not a flood error
        assert limiter._paused_until == 0

    @pytest.mark.asyncio
    async def test_flood_with_seconds_attribute(self):
        """Exception with .seconds attribute uses that value for pause."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        class FloodWaitCustom(Exception):
            def __init__(self):
                self.seconds = 2
                super().__init__("Flood wait custom")

        async def mock_flood():
            raise FloodWaitCustom()

        with contextlib.suppress(Exception):
            await limiter.enqueue(mock_flood, dedup_key="flood_sec")

        assert limiter._paused_until > 0

    @pytest.mark.asyncio
    async def test_proactive_strict_sliding_window(self):
        """
        Proactive limiter should enforce a strict sliding window:
        for any i, t[i+rate_limit] - t[i] >= rate_window (within tolerance).
        """
        limiter = self.create_limiter(rate_limit=2, rate_window=0.5)

        async def acquire(i: int) -> float:
            async def _do() -> float:
                return time.monotonic()

            return await limiter.enqueue(_do, dedup_key=f"strict:{i}")

        acquired = await asyncio.gather(*(acquire(i) for i in range(5)))
        acquired.sort()

        rate_limit = 2
        rate_window = 0.5
        tolerance = 0.05
        for i in range(len(acquired) - rate_limit):
            assert acquired[i + rate_limit] - acquired[i] >= rate_window - tolerance, (
                f"Sliding window violated at i={i}: "
                f"dt={acquired[i + rate_limit] - acquired[i]:.3f}s"
            )

    @pytest.mark.asyncio
    async def test_compaction_last_task_fails_all_futures_get_exception(self):
        """When compacted task's last func fails, all futures get the exception."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        async def ok_task():
            return "ok"

        async def fail_task():
            raise RuntimeError("last task failed")

        future1 = asyncio.create_task(limiter.enqueue(ok_task, dedup_key="fail_key"))
        future2 = asyncio.create_task(limiter.enqueue(fail_task, dedup_key="fail_key"))

        with pytest.raises(RuntimeError, match="last task failed"):
            await future1
        with pytest.raises(RuntimeError, match="last task failed"):
            await future2

    @pytest.mark.asyncio
    async def test_fire_and_forget_failure_logged(self, caplog):
        """fire_and_forget with failing task logs error and does not re-raise."""
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)

        async def fail_task():
            raise ValueError("fire_and_forget failed")

        limiter.fire_and_forget(fail_task, dedup_key="fire_fail")
        await asyncio.sleep(1.5)

        joined = " ".join(str(r.message) for r in caplog.records)
        assert "ValueError" in joined
        assert "fire_and_forget failed" not in joined

    @pytest.mark.asyncio
    async def test_shutdown_settles_active_queued_and_background_work(self):
        limiter = self.create_limiter(rate_limit=1, rate_window=60.0)
        active_started = asyncio.Event()
        never_finish = asyncio.Event()

        async def active_operation() -> None:
            active_started.set()
            await never_finish.wait()

        active = asyncio.create_task(
            limiter.enqueue(active_operation, dedup_key="active")
        )
        await active_started.wait()

        async def queued_operation() -> None:
            await never_finish.wait()

        queued = asyncio.create_task(
            limiter.enqueue(queued_operation, dedup_key="queued")
        )
        limiter.fire_and_forget(queued_operation, dedup_key="background")
        await asyncio.sleep(0)

        await limiter.shutdown(timeout=0.1)
        results = await asyncio.gather(active, queued, return_exceptions=True)

        assert all(isinstance(result, asyncio.CancelledError) for result in results)
        assert limiter._background_tasks == set()
        assert limiter._queue_map == {}
        assert not limiter._queue_list

    @pytest.mark.asyncio
    async def test_enqueue_cannot_enter_after_shutdown_begins(self):
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)
        await limiter._condition.acquire()

        async def succeed() -> str:
            return "unexpected"

        enqueue_task = asyncio.create_task(
            limiter.enqueue(succeed, dedup_key="shutdown-race")
        )
        await asyncio.sleep(0)
        shutdown_task = asyncio.create_task(limiter.shutdown(timeout=0.1))
        await asyncio.sleep(0)
        assert limiter._closed is True

        limiter._condition.release()
        await shutdown_task

        with pytest.raises(RuntimeError, match="is closed"):
            await enqueue_task

    @pytest.mark.asyncio
    async def test_shutdown_preserves_external_cancellation(self):
        limiter = self.create_limiter(rate_limit=1, rate_window=1.0)
        release = asyncio.Event()

        async def cancellation_resistant_worker() -> None:
            try:
                await release.wait()
            except asyncio.CancelledError:
                await release.wait()

        worker_task = limiter._worker_task
        assert worker_task is not None
        await asyncio.sleep(0)
        worker_task.cancel()
        await asyncio.gather(worker_task, return_exceptions=True)
        limiter._worker_task = asyncio.create_task(cancellation_resistant_worker())
        shutdown_task = asyncio.create_task(limiter.shutdown(timeout=1.0))
        await asyncio.sleep(0)

        shutdown_task.cancel()
        release.set()

        with pytest.raises(asyncio.CancelledError):
            await shutdown_task

    @pytest.mark.asyncio
    async def test_cancelled_shutdown_retries_queued_future_settlement(self):
        limiter = self.create_limiter(rate_limit=1, rate_window=60.0)
        active_started = asyncio.Event()
        never_finish = asyncio.Event()

        async def active_operation() -> None:
            active_started.set()
            await never_finish.wait()

        active = asyncio.create_task(
            limiter.enqueue(active_operation, dedup_key="active")
        )
        await active_started.wait()

        async def queued_operation() -> None:
            await never_finish.wait()

        queued = asyncio.create_task(
            limiter.enqueue(queued_operation, dedup_key="queued")
        )
        while "queued" not in limiter._queue_map:
            await asyncio.sleep(0)

        await limiter._condition.acquire()
        shutdown_task = asyncio.create_task(limiter.shutdown())
        await asyncio.sleep(0)
        assert limiter._closed is True

        shutdown_task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await shutdown_task
        limiter._condition.release()

        await limiter.shutdown(timeout=0.1)
        results = await asyncio.gather(active, queued, return_exceptions=True)

        assert all(isinstance(result, asyncio.CancelledError) for result in results)
        assert limiter._queue_map == {}
        assert not limiter._queue_list
        assert limiter._worker_task is None

    @pytest.mark.asyncio
    async def test_cancelled_operation_does_not_stop_owned_worker(self):
        limiter = self.create_limiter(rate_limit=2, rate_window=1.0)

        async def cancelled_operation() -> None:
            raise asyncio.CancelledError

        async def successful_operation() -> str:
            return "delivered"

        first = asyncio.create_task(
            limiter.enqueue(cancelled_operation, dedup_key="cancelled")
        )
        second = asyncio.create_task(
            limiter.enqueue(successful_operation, dedup_key="next")
        )
        results = await asyncio.gather(first, second, return_exceptions=True)

        assert isinstance(results[0], asyncio.CancelledError)
        assert results[1] == "delivered"
        assert limiter._worker_task is not None
        assert not limiter._worker_task.done()