Do not make a real call for the BeforeRequestHook dummy request

Author: awalker4

_test_unstructured_client/integration/test_decorators.py

@@ -291,27 +291,26 @@ async def test_split_pdf_requests_do_retry(monkeypatch):
     """
     Test that when we split a pdf, the split requests will honor retryable errors.
     """
+    mock_endpoint_called = False
     number_of_split_502s = 2
     number_of_last_page_502s = 2
 
-    async def mock_send(_, request):
+    async def mock_send(_, request: httpx.Request):
         """
         Return a predefined number of 502s for requests with certain starting_page_number values.
 
-        This is because N-1 splits are sent off in the hook logic. These need explicit retry handling.
-        The final split is returned to the SDK and gets the built in retry code.
+        This is to make sure specific portions of the doc are retried properly.
 
         We want to make sure both code paths are retried.
         """
-        request_body = request.read()
+        # Assert that the SDK issues our no-op request
+        # returned by the BeforeRequestHook
+        nonlocal mock_endpoint_called
+        if request.url.host == "no-op":
+            mock_endpoint_called = True
+            return Response(200, request=request)
 
-        # Always return a 200 if the request is not for us
-        # (This mocks the httpbin.org call that gets us into the after_success hook)
-        if "/general/v0/general" not in str(request.url):
-            return Response(
-                200,
-                request=request,
-            )
+        request_body = request.read()
 
         decoded_body = MultipartDecoder(request_body, request.headers.get("Content-Type"))
         form_data = form_utils.parse_form_data(decoded_body)
@@ -369,4 +368,6 @@ async def mock_send(_, request):
 
     assert number_of_split_502s == 0
     assert number_of_last_page_502s == 0
+    assert mock_endpoint_called
+
     assert res.status_code == 200

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -5,6 +5,7 @@
 from collections import Counter
 from typing import Coroutine
 
+import httpx
 import pytest
 import requests
 from requests_toolbelt import MultipartDecoder, MultipartEncoder
@@ -32,7 +33,7 @@ def test_unit_sdk_init():
     hook = SplitPdfHook()
     # This is a fake URL, test doesn't make an API call
     test_url = "http://localhost:5000"
-    test_client = requests.Session()
+    test_client = httpx.Client()
 
     hook.sdk_init(test_url, test_client)
 

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -116,15 +116,38 @@ def sdk_init(
     ) -> Tuple[str, HttpClient]:
         """Initializes Split PDF Hook.
 
+        Adds a mock transport layer to the httpx client. This will return an
+        empty 200 response whenever the specified "dummy host" is used. The before_request
+        hook returns this request so the SDK always succeeds and jumps straight to
+        after_success, where we can await the split results.
+
         Args:
             base_url (str): URL of the API.
             client (HttpClient): HTTP Client.
 
         Returns:
-            Tuple[str, httpx.Session]: The initialized SDK options.
+            Tuple[str, HttpClient]: The initialized SDK options.
         """
-        self.client = client
-        return base_url, client
+        class DummyTransport(httpx.BaseTransport):
+            def __init__(self, base_transport: httpx.BaseTransport):
+                self.base_transport = base_transport
+
+            def handle_request(self, request: httpx.Request) -> httpx.Response:
+                # Return an empty 200 response if we send a request to this dummy host
+                if request.method == "GET" and request.url.host == "no-op":
+                    return httpx.Response(status_code=200, content=b'')
+
+                # Otherwise, pass the request to the default transport
+                return self.base_transport.handle_request(request)
+
+        # Explicit cast to httpx.Client to avoid a typing error
+        httpx_client = cast(httpx.Client, client)
+
+        # pylint: disable=protected-access
+        httpx_client._transport = DummyTransport(httpx_client._transport)
+
+        self.client = httpx_client
+        return base_url, self.client
 
     # pylint: disable=too-many-return-statements
     def before_request(
@@ -289,7 +312,7 @@ async def call_api_partial(page):
 
         # Return a dummy request for the SDK to use
         # This allows us to skip right to the AfterRequestHook and await all the calls
-        dummy_request = httpx.Request("GET",  "https://httpbin.org/status/200")
+        dummy_request = httpx.Request("GET",  "http://no-op")
 
         return dummy_request