|
| 1 | +import asyncio |
| 2 | +import io |
| 3 | + |
| 4 | +import httpx |
| 5 | +import pytest |
| 6 | +from pypdf import PdfWriter |
| 7 | + |
| 8 | +from unstructured_client import UnstructuredClient |
| 9 | +from unstructured_client._hooks.custom.split_pdf_hook import _get_collection_url |
| 10 | +from unstructured_client.models import operations, shared |
| 11 | +from unstructured_client.utils import BackoffStrategy, RetryConfig |
| 12 | + |
| 13 | + |
| 14 | +PREFIXES = [ |
| 15 | + "", |
| 16 | + "/deployment", |
| 17 | + "/one/two", |
| 18 | + "/general/v0/general/deployment", |
| 19 | + "/encoded%2Fprefix", |
| 20 | +] |
| 21 | + |
| 22 | + |
| 23 | +@pytest.fixture |
| 24 | +def pdf_bytes(): |
| 25 | + document = PdfWriter() |
| 26 | + for _ in range(4): |
| 27 | + document.add_blank_page(width=612, height=792) |
| 28 | + buffer = io.BytesIO() |
| 29 | + document.write(buffer) |
| 30 | + return buffer.getvalue() |
| 31 | + |
| 32 | + |
| 33 | +@pytest.mark.parametrize("prefix", PREFIXES) |
| 34 | +def test_collection_url_preserves_encoded_prefix_and_replaces_terminal_route(prefix): |
| 35 | + url = httpx.URL( |
| 36 | + f"https://local.invalid{prefix}/general/v0/general?strategy=fast#fragment" |
| 37 | + ) |
| 38 | + assert ( |
| 39 | + str(_get_collection_url(url)) == f"https://local.invalid{prefix}/general/docs" |
| 40 | + ) |
| 41 | + |
| 42 | + |
| 43 | +def make_sdk(monkeypatch, prefix, override): |
| 44 | + sent = [] |
| 45 | + |
| 46 | + def respond(request): |
| 47 | + sent.append((request.method, str(request.url))) |
| 48 | + return httpx.Response(200, json=[{"text": "hello"}], request=request) |
| 49 | + |
| 50 | + transport = httpx.MockTransport(respond) |
| 51 | + original_async_client = httpx.AsyncClient |
| 52 | + sync_client = httpx.Client(transport=transport) |
| 53 | + async_client = original_async_client(transport=transport) |
| 54 | + sdk = UnstructuredClient( |
| 55 | + server_url="http://unused.invalid" |
| 56 | + if override |
| 57 | + else f"http://local.invalid{prefix}", |
| 58 | + client=sync_client, |
| 59 | + async_client=async_client, |
| 60 | + ) |
| 61 | + # Split chunks create their own async HTTP client; keep every request offline. |
| 62 | + monkeypatch.setattr( |
| 63 | + httpx, |
| 64 | + "AsyncClient", |
| 65 | + lambda **kwargs: original_async_client(transport=transport, **kwargs), |
| 66 | + ) |
| 67 | + return sdk, sent, sync_client, async_client |
| 68 | + |
| 69 | + |
| 70 | +def request(pdf_bytes, split): |
| 71 | + return operations.PartitionRequest( |
| 72 | + partition_parameters=shared.PartitionParameters( |
| 73 | + files=shared.Files(content=pdf_bytes, file_name="test.pdf"), |
| 74 | + split_pdf_page=split, |
| 75 | + ) |
| 76 | + ) |
| 77 | + |
| 78 | + |
| 79 | +def assert_routes(sent, prefix, split): |
| 80 | + partition = f"http://local.invalid{prefix}/general/v0/general" |
| 81 | + if split: |
| 82 | + assert sent[0] == ("GET", f"http://local.invalid{prefix}/general/docs") |
| 83 | + assert len(sent) == 3 |
| 84 | + assert sent[1:] == [("POST", partition)] * 2 |
| 85 | + else: |
| 86 | + assert sent == [("POST", partition)] |
| 87 | + |
| 88 | + |
| 89 | +@pytest.mark.parametrize("prefix", PREFIXES) |
| 90 | +@pytest.mark.parametrize("override", [False, True]) |
| 91 | +@pytest.mark.parametrize("split", [False, True]) |
| 92 | +def test_public_sync_sdk_preserves_deployment_routes( |
| 93 | + monkeypatch, pdf_bytes, prefix, override, split |
| 94 | +): |
| 95 | + sdk, sent, sync_client, async_client = make_sdk(monkeypatch, prefix, override) |
| 96 | + result = sdk.general.partition( |
| 97 | + request=request(pdf_bytes, split), |
| 98 | + server_url=f"http://local.invalid{prefix}" if override else None, |
| 99 | + retries=RetryConfig("none", BackoffStrategy(1, 1, 1, 1), False), |
| 100 | + ) |
| 101 | + assert result.status_code == 200 |
| 102 | + assert_routes(sent, prefix, split) |
| 103 | + sync_client.close() |
| 104 | + asyncio.run(async_client.aclose()) |
| 105 | + |
| 106 | + |
| 107 | +@pytest.mark.asyncio |
| 108 | +@pytest.mark.parametrize("prefix", PREFIXES) |
| 109 | +@pytest.mark.parametrize("override", [False, True]) |
| 110 | +@pytest.mark.parametrize("split", [False, True]) |
| 111 | +async def test_public_async_sdk_preserves_deployment_routes( |
| 112 | + monkeypatch, pdf_bytes, prefix, override, split |
| 113 | +): |
| 114 | + sdk, sent, sync_client, async_client = make_sdk(monkeypatch, prefix, override) |
| 115 | + result = await sdk.general.partition_async( |
| 116 | + request=request(pdf_bytes, split), |
| 117 | + server_url=f"http://local.invalid{prefix}" if override else None, |
| 118 | + retries=RetryConfig("none", BackoffStrategy(1, 1, 1, 1), False), |
| 119 | + ) |
| 120 | + assert result.status_code == 200 |
| 121 | + assert_routes(sent, prefix, split) |
| 122 | + sync_client.close() |
| 123 | + await async_client.aclose() |
0 commit comments