Add the no-op client mock for httpx.AsyncClient

The hooks do not give us access to the AsyncClient, which breaks the mock no-op request in the
splitting hook. For now, let's add it as a param to the sdk_init hook. This means we'll have to
bring this code back whenever the SDK regenerates. At least, until we have an update from Speakeasy
for this.

Author: awalker4

_test_unstructured_client/integration/test_integration_freemium.py

@@ -101,6 +101,29 @@ def test_partition_handling_server_error(error, split_pdf, monkeypatch, doc_path
         response = client.general.partition(request=req)
 
 
+@pytest.mark.asyncio
+async def test_partition_async_returns_elements(client, doc_path):
+    filename = "layout-parser-paper.pdf"
+    with open(doc_path / filename, "rb") as f:
+        files = shared.Files(
+            content=f.read(),
+            file_name=filename,
+        )
+
+    req = operations.PartitionRequest(
+        partition_parameters=shared.PartitionParameters(
+            files=files,
+            strategy="fast",
+            languages=["eng"],
+            split_pdf_page=True,
+        )
+    )
+
+    response = await client.general.partition_async(request=req)
+    assert response.status_code == 200
+    assert len(response.elements)
+
+
 def test_uvloop_partitions_without_errors(client, doc_path):
     async def call_api():
         filename = "layout-parser-paper-fast.pdf"

src/unstructured_client/_hooks/custom/clean_server_url_hook.py

@@ -2,7 +2,7 @@
 from urllib.parse import ParseResult, urlparse, urlunparse
 
 from unstructured_client._hooks.types import SDKInitHook
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import HttpClient, AsyncHttpClient
 
 
 class CleanServerUrlSDKInitHook(SDKInitHook):
@@ -25,9 +25,9 @@ def clean_server_url(self, base_url) -> str:
         return urlunparse(parsed_url._replace(path=""))
 
     def sdk_init(
-        self, base_url: str, client: HttpClient
-    ) -> Tuple[str, HttpClient]:
+        self, base_url: str, client: HttpClient, async_client: AsyncHttpClient
+    ) -> Tuple[str, HttpClient, AsyncHttpClient]:
         """Concrete implementation for SDKInitHook."""
         cleaned_url = self.clean_server_url(base_url)
 
-        return cleaned_url, client
+        return cleaned_url, client, async_client

src/unstructured_client/_hooks/custom/logger_hook.py

@@ -11,7 +11,7 @@
     SDKInitHook,
     AfterSuccessHook,
 )
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import HttpClient, AsyncHttpClient
 from collections import defaultdict
 
 logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
@@ -46,10 +46,10 @@ def log_retries(self, response: Optional[httpx.Response],  error: Optional[Excep
 
 
     def sdk_init(
-        self, base_url: str, client: HttpClient
-    ) -> Tuple[str, HttpClient]:
+        self, base_url: str, client: HttpClient, async_client: AsyncHttpClient
+    ) -> Tuple[str, HttpClient, AsyncHttpClient]:
         logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
-        return base_url, client
+        return base_url, client, async_client
 
     def after_success(
         self, hook_ctx: AfterSuccessContext, response: httpx.Response

src/unstructured_client/_hooks/custom/split_pdf_hook.py

@@ -32,7 +32,7 @@
     BeforeRequestHook,
     SDKInitHook,
 )
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import HttpClient, AsyncHttpClient
 from unstructured_client.models import shared
 
 logger = logging.getLogger(UNSTRUCTURED_CLIENT_LOGGER_NAME)
@@ -104,6 +104,7 @@ class SplitPdfHook(SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorH
 
     def __init__(self) -> None:
         self.client: Optional[HttpClient] = None
+        self.async_client: Optional[AsyncHttpClient] = None
         self.coroutines_to_execute: dict[
             str, list[Coroutine[Any, Any, httpx.Response]]
         ] = {}
@@ -112,8 +113,8 @@ def __init__(self) -> None:
         self.allow_failed: bool = DEFAULT_ALLOW_FAILED
 
     def sdk_init(
-            self, base_url: str, client: HttpClient
-    ) -> Tuple[str, HttpClient]:
+            self, base_url: str, client: HttpClient, async_client: AsyncHttpClient
+    ) -> Tuple[str, HttpClient, AsyncHttpClient]:
         """Initializes Split PDF Hook.
 
         Adds a mock transport layer to the httpx client. This will return an
@@ -140,14 +141,31 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
                 # Otherwise, pass the request to the default transport
                 return self.base_transport.handle_request(request)
 
+        class AsyncDummyTransport(httpx.AsyncBaseTransport):
+            def __init__(self, base_transport: httpx.AsyncBaseTransport):
+                self.base_transport = base_transport
+
+            async def handle_async_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 await self.base_transport.handle_async_request(request)
+
         # Explicit cast to httpx.Client to avoid a typing error
         httpx_client = cast(httpx.Client, client)
+        async_httpx_client = cast(httpx.AsyncClient, async_client)
 
         # pylint: disable=protected-access
         httpx_client._transport = DummyTransport(httpx_client._transport)
 
+        # pylint: disable=protected-access
+        async_httpx_client._transport = AsyncDummyTransport(async_httpx_client._transport)
+
         self.client = httpx_client
-        return base_url, self.client
+        self.async_client = async_httpx_client
+        return base_url, self.client, self.async_client
 
     # pylint: disable=too-many-return-statements
     def before_request(

src/unstructured_client/_hooks/sdkhooks.py

@@ -4,7 +4,7 @@
 from .types import SDKInitHook, BeforeRequestContext, BeforeRequestHook, AfterSuccessContext, AfterSuccessHook, AfterErrorContext, AfterErrorHook, Hooks
 from .registration import init_hooks
 from typing import List, Optional, Tuple
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import HttpClient, AsyncHttpClient
 
 class SDKHooks(Hooks):
     def __init__(self) -> None:
@@ -26,10 +26,10 @@ def register_after_success_hook(self, hook: AfterSuccessHook) -> None:
     def register_after_error_hook(self, hook: AfterErrorHook) -> None:
         self.after_error_hooks.append(hook)
 
-    def sdk_init(self, base_url: str, client: HttpClient) -> Tuple[str, HttpClient]:
+    def sdk_init(self, base_url: str, client: HttpClient, async_client: AsyncHttpClient) -> Tuple[str, HttpClient, AsyncHttpClient]:
         for hook in self.sdk_init_hooks:
-            base_url, client = hook.sdk_init(base_url, client)
-        return base_url, client
+            base_url, client, async_client = hook.sdk_init(base_url, client, async_client)
+        return base_url, client, async_client
 
     def before_request(self, hook_ctx: BeforeRequestContext, request: httpx.Request) -> httpx.Request:
         for hook in self.before_request_hooks:

src/unstructured_client/_hooks/types.py

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 import httpx
 from typing import Any, Callable, List, Optional, Tuple, Union
-from unstructured_client.httpclient import HttpClient
+from unstructured_client.httpclient import HttpClient, AsyncHttpClient
 
 
 class HookContext:
@@ -36,7 +36,7 @@ def __init__(self, hook_ctx: HookContext):
 
 class SDKInitHook(ABC):
     @abstractmethod
-    def sdk_init(self, base_url: str, client: HttpClient) -> Tuple[str, HttpClient]:
+    def sdk_init(self, base_url: str, client: HttpClient, async_client: AsyncHttpClient) -> Tuple[str, HttpClient, AsyncHttpClient]:
         pass
 
 

src/unstructured_client/general.py

@@ -94,8 +94,9 @@ def partition(
 
 
     async def partition_async(
-        self, *,
+        self,
         request: Union[operations.PartitionRequest, operations.PartitionRequestTypedDict],
+        *,
         retries: OptionalNullable[utils.RetryConfig] = UNSET,
         server_url: Optional[str] = None,
         timeout_ms: Optional[int] = None,

src/unstructured_client/sdk.py

@@ -74,7 +74,7 @@ def __init__(
         hooks = SDKHooks()
 
         current_server_url, *_ = self.sdk_configuration.get_server_details()
-        server_url, self.sdk_configuration.client = hooks.sdk_init(current_server_url, self.sdk_configuration.client)
+        server_url, self.sdk_configuration.client, self.sdk_configuration.async_client = hooks.sdk_init(current_server_url, self.sdk_configuration.client, self.sdk_configuration.async_client)
         if current_server_url != server_url:
             self.sdk_configuration.server_url = server_url