|
4 | 4 | import json |
5 | 5 | from dataclasses import replace |
6 | 6 | from types import SimpleNamespace |
| 7 | +from typing import TYPE_CHECKING |
7 | 8 |
|
8 | 9 | import httpx |
9 | 10 | import pytest |
|
17 | 18 | from deepagents_talon.history_vector_backends import vector_backend |
18 | 19 | from tests.unit_tests.test_history_vectors import SCOPE, settled |
19 | 20 |
|
| 21 | +if TYPE_CHECKING: |
| 22 | + from collections.abc import Awaitable |
| 23 | + from pathlib import Path |
| 24 | + |
20 | 25 | PREFIX = "DEEPAGENTS_TALON_HISTORY_EMBED_" |
21 | 26 |
|
22 | 27 |
|
@@ -519,43 +524,57 @@ def driver(module, _extra): |
519 | 524 | assert type(index["embed"]).__name__ == "AutoEmbeddings" |
520 | 525 |
|
521 | 526 |
|
522 | | -async def test_rate_limit_backoff_survives_new_appends(tmp_path, monkeypatch): |
523 | | - from deepagents_talon import history_vectors # noqa: PLC0415 |
| 527 | +async def test_rate_limit_backoff_survives_new_appends( |
| 528 | + tmp_path: Path, monkeypatch: pytest.MonkeyPatch |
| 529 | +) -> None: |
| 530 | + now = 0.0 |
| 531 | + waits: asyncio.Queue[float] = asyncio.Queue() |
| 532 | + |
| 533 | + async def wait_for(wake: Awaitable[object], timeout: float) -> None: # noqa: ASYNC109 # Match asyncio.wait_for's interface. |
| 534 | + waits.put_nowait(timeout) |
| 535 | + await asyncio.wait_for(wake, 5) |
| 536 | + |
| 537 | + # Control only the worker's clock; real asyncio deadlines still bound the test. |
| 538 | + worker_asyncio = SimpleNamespace(**vars(asyncio)) |
| 539 | + worker_asyncio.get_running_loop = lambda: SimpleNamespace(time=lambda: now) |
| 540 | + worker_asyncio.wait_for = wait_for |
| 541 | + monkeypatch.setattr(history_vectors, "asyncio", worker_asyncio) |
524 | 542 |
|
525 | 543 | class RateLimited(RecordingEmbeddings): |
526 | 544 | attempts = 0 |
527 | | - failed = asyncio.Event() |
528 | | - succeeded = asyncio.Event() |
529 | 545 |
|
530 | | - async def aembed_documents(self, texts): |
| 546 | + async def aembed_documents(self, texts: list[str]) -> list[list[float]]: |
531 | 547 | self.attempts += 1 |
532 | 548 | if self.attempts == 1: |
533 | | - self.failed.set() |
534 | | - response = httpx.Response(429, headers={"retry-after": "0.1"}) |
| 549 | + response = httpx.Response(429, headers={"retry-after": "120"}) |
535 | 550 | msg = "rate limited" |
536 | 551 | raise httpx.HTTPStatusError( |
537 | 552 | msg, |
538 | 553 | request=httpx.Request("POST", "https://test.invalid"), |
539 | 554 | response=response, |
540 | 555 | ) |
541 | | - self.succeeded.set() |
542 | 556 | return await super().aembed_documents(texts) |
543 | 557 |
|
544 | 558 | raw = RateLimited() |
545 | 559 | fake_adapter(monkeypatch, raw) |
546 | | - monkeypatch.setattr(history_vectors, "_RETRY_SECONDS", 0.001) |
547 | 560 | monkeypatch.setattr(history_vectors.secrets, "randbelow", lambda _bound: 0) |
548 | 561 | config = configuration(tmp_path) |
549 | 562 | async with open_history(config) as archive: |
| 563 | + assert await asyncio.wait_for(waits.get(), 5) == history_vectors._RETRY_SECONDS |
550 | 564 | await archive.append(SCOPE, "first", "time", [HumanMessage("first")]) |
551 | | - await asyncio.wait_for(raw.failed.wait(), 1) |
| 565 | + assert await asyncio.wait_for(waits.get(), 5) == 120 |
552 | 566 | for i in range(3): |
553 | | - await archive.append(SCOPE, f"later-{i}", "time", [HumanMessage("later")]) |
554 | | - assert not raw.succeeded.is_set() |
555 | | - assert await archive.vectors.archive.pending(SCOPE) |
556 | | - await asyncio.wait_for(raw.succeeded.wait(), 1) |
| 567 | + now += 20 |
| 568 | + await archive.append(SCOPE, f"later-{i}", "time", [HumanMessage(f"later-{i}")]) |
| 569 | + assert await asyncio.wait_for(waits.get(), 5) == 120 - now |
| 570 | + assert raw.attempts == 1 |
| 571 | + assert await archive.vectors.archive.pending(SCOPE) |
| 572 | + now = 120 |
| 573 | + archive.vectors.wake.set() |
557 | 574 | await settled(archive) |
558 | 575 | assert not await archive.vectors.archive.pending(SCOPE) |
| 576 | + stored = await archive.vectors.store.asearch(archive.vectors.namespace("whatsapp", "one")) |
| 577 | + assert {item.value["text"] for item in stored} == {"first", "later-0", "later-1", "later-2"} |
559 | 578 |
|
560 | 579 |
|
561 | 580 | async def test_threaded_store_uses_native_async_embeddings(tmp_path): |
|
0 commit comments