|
5 | 5 | import httpx |
6 | 6 | import pytest |
7 | 7 |
|
| 8 | +from unstructured_client import UnstructuredClient |
| 9 | +from unstructured_client.models.errors import ServerError |
8 | 10 | from unstructured_client.utils import retries as retries_module |
9 | 11 | from unstructured_client.utils.retries import ( |
10 | 12 | BackoffStrategy, |
@@ -555,3 +557,161 @@ async def func(): |
555 | 557 | with pytest.raises(ValueError, match="permanent"): |
556 | 558 | asyncio.run(retry_async(func, retries_config)) |
557 | 559 | assert call_count == 1 |
| 560 | + |
| 561 | + |
| 562 | +def _send_with_retry(respond, config, async_mode, *, stream=False): |
| 563 | + transport = httpx.MockTransport(respond) |
| 564 | + if async_mode: |
| 565 | + async def run(): |
| 566 | + async with httpx.AsyncClient(transport=transport) as client: |
| 567 | + return await retry_async( |
| 568 | + lambda: client.send(client.build_request("GET", "http://test.invalid"), stream=stream), |
| 569 | + config, |
| 570 | + ) |
| 571 | + return asyncio.run(run()) |
| 572 | + with httpx.Client(transport=transport) as client: |
| 573 | + return retry( |
| 574 | + lambda: client.send(client.build_request("GET", "http://test.invalid"), stream=stream), |
| 575 | + config, |
| 576 | + ) |
| 577 | + |
| 578 | + |
| 579 | +@pytest.mark.parametrize("async_mode", [False, True]) |
| 580 | +@pytest.mark.parametrize( |
| 581 | + "body,content_type,stop", |
| 582 | + [ |
| 583 | + (b'{"detail":"unsupported","retryable":false}', "application/json", True), |
| 584 | + (b'{"retryable":false}', "application/problem+json; charset=utf-8", True), |
| 585 | + (b'{"retryable":true}', "application/json", False), |
| 586 | + (b'{"detail":"Internal server error"}', "application/json", False), |
| 587 | + (b'{"retryable":"false"}', "application/json", False), |
| 588 | + (b'{"retryable":0}', "application/json", False), |
| 589 | + (b'{"retryable":null}', "application/json", False), |
| 590 | + (b'[{"retryable":false}]', "application/json", False), |
| 591 | + (b'false', "application/json", False), |
| 592 | + (b'{broken', "application/json", False), |
| 593 | + pytest.param( |
| 594 | + b'{"retryable":false,"detail":' + b"[" * 100_000 + b"]" * 100_000 + b"}", |
| 595 | + "application/json", |
| 596 | + False, |
| 597 | + id="deeply-nested", |
| 598 | + ), |
| 599 | + (b'upstream error', "text/plain", False), |
| 600 | + (b'{"retryable":false}', "text/plain", True), |
| 601 | + (b'{"retryable":false}', "", True), |
| 602 | + ], |
| 603 | +) |
| 604 | +def test_declared_retry_verdict( |
| 605 | + fake_clock, async_mode, body, content_type, stop |
| 606 | +): |
| 607 | + status_code = 503 |
| 608 | + attempts = [] |
| 609 | + |
| 610 | + def respond(request): |
| 611 | + attempts.append(request) |
| 612 | + if len(attempts) == 1: |
| 613 | + return httpx.Response(status_code, content=body, headers={"Content-Type": content_type}) |
| 614 | + return httpx.Response(200, json=[]) |
| 615 | + |
| 616 | + config = _retries(status_codes=["5xx"], initial_interval=1, max_interval=1) |
| 617 | + result = _send_with_retry(respond, config, async_mode) |
| 618 | + |
| 619 | + assert len(attempts) == (1 if stop else 2) |
| 620 | + assert result.status_code == (status_code if stop else 200) |
| 621 | + if stop: |
| 622 | + assert result.content == body |
| 623 | + |
| 624 | + |
| 625 | +@pytest.mark.parametrize("async_mode", [False, True]) |
| 626 | +@pytest.mark.parametrize("verdict", [False, True]) |
| 627 | +def test_declared_verdict_does_not_retry_non_matching_status(fake_clock, async_mode, verdict): |
| 628 | + attempts = [] |
| 629 | + |
| 630 | + def respond(request): |
| 631 | + attempts.append(request) |
| 632 | + return httpx.Response(400, json={"retryable": verdict}) |
| 633 | + |
| 634 | + config = _retries(status_codes=["5xx"]) |
| 635 | + response = _send_with_retry(respond, config, async_mode) |
| 636 | + |
| 637 | + assert response.status_code == 400 |
| 638 | + assert len(attempts) == 1 |
| 639 | + |
| 640 | + |
| 641 | +@pytest.mark.parametrize("async_mode", [False, True]) |
| 642 | +def test_unread_stream_is_not_consumed(fake_clock, async_mode): |
| 643 | + class SyncStream(httpx.SyncByteStream): |
| 644 | + def __iter__(self): |
| 645 | + pytest.fail("retry inspected an unread stream") |
| 646 | + yield b"" |
| 647 | + |
| 648 | + class AsyncStream(httpx.AsyncByteStream): |
| 649 | + async def __aiter__(self): |
| 650 | + pytest.fail("retry inspected an unread stream") |
| 651 | + yield b"" |
| 652 | + |
| 653 | + attempts = [] |
| 654 | + |
| 655 | + def respond(request): |
| 656 | + attempts.append(request) |
| 657 | + if len(attempts) == 1: |
| 658 | + return httpx.Response( |
| 659 | + 503, headers={"Content-Type": "application/json"}, |
| 660 | + stream=AsyncStream() if async_mode else SyncStream(), |
| 661 | + ) |
| 662 | + return httpx.Response(200) |
| 663 | + |
| 664 | + config = _retries(status_codes=["5xx"]) |
| 665 | + result = _send_with_retry(respond, config, async_mode, stream=True) |
| 666 | + |
| 667 | + assert result.status_code == 200 |
| 668 | + assert len(attempts) == 2 |
| 669 | + |
| 670 | + |
| 671 | +@pytest.mark.parametrize("async_mode", [False, True]) |
| 672 | +@pytest.mark.parametrize("status_code", [500, 503]) |
| 673 | +def test_public_partition_preserves_declared_error(fake_clock, async_mode, status_code): |
| 674 | + body = {"detail": "unsupported document", "retryable": False} |
| 675 | + attempts = [] |
| 676 | + |
| 677 | + def respond(request): |
| 678 | + attempts.append(request) |
| 679 | + return httpx.Response(status_code, json=body) |
| 680 | + |
| 681 | + async def run(): |
| 682 | + transport = httpx.MockTransport(respond) |
| 683 | + with httpx.Client(transport=transport) as client: |
| 684 | + async with httpx.AsyncClient(transport=transport) as async_client: |
| 685 | + sdk = UnstructuredClient( |
| 686 | + server_url="http://test.invalid", client=client, async_client=async_client |
| 687 | + ) |
| 688 | + request = {"partition_parameters": { |
| 689 | + "files": {"content": b"hello", "file_name": "test.txt"}, |
| 690 | + "split_pdf_page": False, |
| 691 | + }} |
| 692 | + config = _retries(initial_interval=1, max_interval=1, max_elapsed_time=2).config |
| 693 | + with pytest.raises(ServerError) as exc: |
| 694 | + if async_mode: |
| 695 | + await sdk.general.partition_async(request=request, retries=config) |
| 696 | + else: |
| 697 | + sdk.general.partition(request=request, retries=config) |
| 698 | + assert exc.value.raw_response.json() == body |
| 699 | + assert "unsupported document" in str(exc.value) |
| 700 | + |
| 701 | + asyncio.run(run()) |
| 702 | + assert len(attempts) == 1 |
| 703 | + |
| 704 | + |
| 705 | +@pytest.mark.parametrize("async_mode", [False, True]) |
| 706 | +@pytest.mark.parametrize("status_code", [500, 503]) |
| 707 | +@pytest.mark.parametrize("status_codes", [["5xx"], ["500", "503"]]) |
| 708 | +def test_declared_verdict_stops_matching_status(fake_clock, async_mode, status_code, status_codes): |
| 709 | + attempts = [] |
| 710 | + |
| 711 | + def respond(request): |
| 712 | + attempts.append(request) |
| 713 | + return httpx.Response(status_code, json={"retryable": False}) |
| 714 | + |
| 715 | + response = _send_with_retry(respond, _retries(status_codes=status_codes), async_mode) |
| 716 | + assert response.status_code == status_code |
| 717 | + assert len(attempts) == 1 |
0 commit comments