Jack1808/Claude_Code

1import asyncio2import os3import time4 5import pytest6import pytest_asyncio7 8# Set environment variables relative to test execution9os.environ["MESSAGING_RATE_LIMIT"] = "1"10os.environ["MESSAGING_RATE_WINDOW"] = "0.5"11 12import contextlib13 14from messaging.limiter import MessagingRateLimiter15 16 17class TestMessagingRateLimiter:18    """Tests for MessagingRateLimiter."""19 20    @pytest_asyncio.fixture(autouse=True)21    async def reset_limiter(self):22        """Reset singleton and environment before each test."""23        # Ensure the singleton worker is stopped between tests to avoid dangling tasks.24        await MessagingRateLimiter.shutdown_instance(timeout=0.1)25        os.environ["MESSAGING_RATE_LIMIT"] = "1"26        os.environ["MESSAGING_RATE_WINDOW"] = "0.5"27 28        yield29 30        await MessagingRateLimiter.shutdown_instance(timeout=0.1)31 32    @pytest.mark.asyncio33    async def test_singleton_pattern(self):34        """Test that get_instance returns the same object."""35        limiter1 = await MessagingRateLimiter.get_instance()36        limiter2 = await MessagingRateLimiter.get_instance()37        assert limiter1 is limiter238 39    @pytest.mark.asyncio40    async def test_compaction(self):41        """42        Verify multiple rapid requests with same dedup_key are compacted.43        Logic ported from verify_limiter.py44        """45        # Set slow rate for testing compaction46        os.environ["MESSAGING_RATE_LIMIT"] = "1"47        os.environ["MESSAGING_RATE_WINDOW"] = "1.0"48 49        # Must reset instance to pick up new env vars50        MessagingRateLimiter._instance = None51        limiter = await MessagingRateLimiter.get_instance()52 53        call_counts = {}54 55        async def mock_edit(msg_id, content):56            call_counts[msg_id] = call_counts.get(msg_id, 0) + 157            return f"done_{content}"58 59        # Spam 5 edits60        for i in range(5):61            limiter.fire_and_forget(62                lambda i=i: mock_edit("msg1", f"update_{i}"), dedup_key="edit:msg1"63            )64 65        # Wait for processing66        # 1st might go through immediately, subsequent ones queue and compact67        await asyncio.sleep(2.5)68 69        # Expected: ~2 calls (first and last)70        assert call_counts["msg1"] <= 2, (71            f"Expected compaction to reduce calls, but got {call_counts.get('msg1', 0)}"72        )73        assert call_counts["msg1"] >= 1, "Expected at least one call"74 75    @pytest.mark.asyncio76    async def test_compaction_and_futures_resolution(self):77        """78        Verify that even when compacted, all futures resolve to the result of the LAST execution.79        Logic ported from verify_limiter_v2.py80        """81        os.environ["MESSAGING_RATE_LIMIT"] = "1"82        os.environ["MESSAGING_RATE_WINDOW"] = "0.5"83        MessagingRateLimiter._instance = None84        limiter = await MessagingRateLimiter.get_instance()85 86        call_counts = {}87        msg_id = "test_msg_hang"88 89        async def mock_edit(mid, content):90            call_counts[mid] = call_counts.get(mid, 0) + 191            await asyncio.sleep(0.05)92            return f"result_{content}"93 94        async def task(i):95            return await limiter.enqueue(96                lambda i=i: mock_edit(msg_id, f"v{i}"), dedup_key=f"edit:{msg_id}"97            )98 99        start_time = time.time()100 101        # Enqueue 3 tasks concurrently102        results = await asyncio.gather(task(1), task(2), task(3))103 104        duration = time.time() - start_time105 106        # All results should be the LAST one executed107        for res in results:108            assert res == "result_v3", f"Expected result_v3, got {res}"109 110        # Should be reasonably fast111        assert duration < 2.0, "Execution took too long"112 113        # Calls should be compacted114        assert call_counts[msg_id] <= 2, f"Too many actual calls: {call_counts[msg_id]}"115 116    @pytest.mark.asyncio117    async def test_flood_wait_handling(self):118        """Test that FloodWait exceptions pause the worker."""119        MessagingRateLimiter._instance = None120        limiter = await MessagingRateLimiter.get_instance()121 122        # Mock exception with .seconds attribute123        class FloodWait(Exception):124            def __init__(self, seconds):125                self.seconds = seconds126                super().__init__(f"Flood wait {seconds}s")127 128        call_count = 0129 130        async def mock_fail():131            nonlocal call_count132            call_count += 1133            raise FloodWait(1)  # 1 second wait134 135        async def mock_success():136            nonlocal call_count137            call_count += 1138            return "success"139 140        # First call fails and triggers pause141        with contextlib.suppress(Exception):142            await limiter.enqueue(mock_fail, dedup_key="key1")143 144        assert limiter._paused_until > 0145 146        # Enqueue success, it should wait147        start = time.time()148        await limiter.enqueue(mock_success, dedup_key="key2")149        duration = time.time() - start150 151        # Should have waited at least ~1s152        assert duration >= 0.9, (153            f"Should have waited for FloodWait, but took {duration:.2f}s"154        )155        assert call_count == 2156 157    @pytest.mark.asyncio158    async def test_flood_wait_retry_after_parsing(self):159        """Error message with 'retry after N' parses the wait seconds."""160        MessagingRateLimiter._instance = None161        limiter = await MessagingRateLimiter.get_instance()162 163        async def mock_flood():164            raise Exception("Flood wait: retry after 2 seconds")165 166        with contextlib.suppress(Exception):167            await limiter.enqueue(mock_flood, dedup_key="retry_parse")168 169        # Should have parsed "after 2" -> 2 seconds170        assert limiter._paused_until > 0171 172    @pytest.mark.asyncio173    async def test_non_flood_exception_no_pause(self):174        """Non-flood exception doesn't trigger pause."""175        MessagingRateLimiter._instance = None176        limiter = await MessagingRateLimiter.get_instance()177 178        async def mock_error():179            raise ValueError("some regular error")180 181        with contextlib.suppress(ValueError):182            await limiter.enqueue(mock_error, dedup_key="non_flood")183 184        # Should NOT have paused since it's not a flood error185        assert limiter._paused_until == 0186 187    @pytest.mark.asyncio188    async def test_flood_with_seconds_attribute(self):189        """Exception with .seconds attribute uses that value for pause."""190        MessagingRateLimiter._instance = None191        limiter = await MessagingRateLimiter.get_instance()192 193        class FloodWaitCustom(Exception):194            def __init__(self):195                self.seconds = 2196                super().__init__("Flood wait custom")197 198        async def mock_flood():199            raise FloodWaitCustom()200 201        with contextlib.suppress(Exception):202            await limiter.enqueue(mock_flood, dedup_key="flood_sec")203 204        assert limiter._paused_until > 0205 206    @pytest.mark.asyncio207    async def test_proactive_strict_sliding_window(self):208        """209        Proactive limiter should enforce a strict sliding window:210        for any i, t[i+rate_limit] - t[i] >= rate_window (within tolerance).211        """212        os.environ["MESSAGING_RATE_LIMIT"] = "2"213        os.environ["MESSAGING_RATE_WINDOW"] = "0.5"214        MessagingRateLimiter._instance = None215        limiter = await MessagingRateLimiter.get_instance()216 217        async def acquire(i: int) -> float:218            async def _do() -> float:219                return time.monotonic()220 221            return await limiter.enqueue(_do, dedup_key=f"strict:{i}")222 223        acquired = await asyncio.gather(*(acquire(i) for i in range(5)))224        acquired.sort()225 226        rate_limit = 2227        rate_window = 0.5228        tolerance = 0.05229        for i in range(len(acquired) - rate_limit):230            assert acquired[i + rate_limit] - acquired[i] >= rate_window - tolerance, (231                f"Sliding window violated at i={i}: "232                f"dt={acquired[i + rate_limit] - acquired[i]:.3f}s"233            )234 235    @pytest.mark.asyncio236    async def test_compaction_last_task_fails_all_futures_get_exception(self):237        """When compacted task's last func fails, all futures get the exception."""238        MessagingRateLimiter._instance = None239        limiter = await MessagingRateLimiter.get_instance()240 241        async def ok_task():242            return "ok"243 244        async def fail_task():245            raise RuntimeError("last task failed")246 247        future1 = asyncio.create_task(limiter.enqueue(ok_task, dedup_key="fail_key"))248        future2 = asyncio.create_task(limiter.enqueue(fail_task, dedup_key="fail_key"))249 250        with pytest.raises(RuntimeError, match="last task failed"):251            await future1252        with pytest.raises(RuntimeError, match="last task failed"):253            await future2254 255    @pytest.mark.asyncio256    async def test_fire_and_forget_failure_logged(self, caplog):257        """fire_and_forget with failing task logs error and does not re-raise."""258        MessagingRateLimiter._instance = None259        limiter = await MessagingRateLimiter.get_instance()260 261        async def fail_task():262            raise ValueError("fire_and_forget failed")263 264        limiter.fire_and_forget(fail_task, dedup_key="fire_fail")265        await asyncio.sleep(1.5)266 267        assert any("fire_and_forget failed" in str(r) for r in caplog.records)268