Mock sync Weaviate cloud checks

Author: cmpadden

libraries/dagster-weaviate/dagster_weaviate_tests/test_resource.py

@@ -104,18 +104,14 @@ def query_local_weaviate_asset(weaviate_local: WeaviateResource):
 
 # Required to avoid an infinite recursion
 # (since our mock needs to forward some requests to the original httpx.Send)
-real_httpx_send_method = httpx.AsyncClient.send
+real_httpx_async_send_method = httpx.AsyncClient.send
+real_httpx_send_method = httpx.Client.send
 
 
 # this method only returns a mock response if the URL accessed is weaviate.
 # Otherwise, it passes the request through to real_httpx_send_method
 # Important since Weaviate might make some PyPI requests (which don't need to be mocked).
-async def mock_httpx_send_method(self, *args, **kwargs):
-    request = args[0]
-
-    if "weaviate.cloud" not in str(request.url):
-        return await real_httpx_send_method(self, *args, **kwargs)
-
+def _mock_weaviate_cloud_response(request):
     assert WCD_URL in str(request.url)
     assert ("x-cohere-api-key", COHERE_APIKEY) in request.headers.items()
     assert WCD_APIKEY in request.headers["authorization"]
@@ -127,7 +123,7 @@ async def mock_httpx_send_method(self, *args, **kwargs):
     #   -H "x-weaviate-api-key: $WCD_API_KEY" \
     #   -H "x-weaviate-cluster-url: $WCD_URL" \
     #   -H "authorization: Bearer $WCD_API_KEY"
-    result = httpx.Response(
+    return httpx.Response(
         status_code=200,
         text=json.dumps(
             {
@@ -156,7 +152,23 @@ async def mock_httpx_send_method(self, *args, **kwargs):
         ),
     )
 
-    return result
+
+async def mock_httpx_async_send_method(self, *args, **kwargs):
+    request = args[0]
+
+    if "weaviate.cloud" not in str(request.url):
+        return await real_httpx_async_send_method(self, *args, **kwargs)
+
+    return _mock_weaviate_cloud_response(request)
+
+
+def mock_httpx_send_method(self, *args, **kwargs):
+    request = kwargs.get("request") or args[0]
+
+    if "weaviate.cloud" not in str(request.url):
+        return real_httpx_send_method(self, *args, **kwargs)
+
+    return _mock_weaviate_cloud_response(request)
 
 
 def test_cloud_resource():
@@ -165,8 +177,8 @@ def test_cloud_resource():
     specified in the Resource config.
     """
     with mock.patch.object(
-        httpx.AsyncClient, "send", autospec=True, side_effect=mock_httpx_send_method
-    ):
+        httpx.AsyncClient, "send", autospec=True, side_effect=mock_httpx_async_send_method
+    ), mock.patch.object(httpx.Client, "send", autospec=True, side_effect=mock_httpx_send_method):
 
         @asset
         def query_cloud_weaviate_asset(weaviate_cloud: WeaviateResource):