@@ -104,18 +104,14 @@ def query_local_weaviate_asset(weaviate_local: WeaviateResource):
104104
105105# Required to avoid an infinite recursion
106106# (since our mock needs to forward some requests to the original httpx.Send)
107- real_httpx_send_method = httpx .AsyncClient .send
107+ real_httpx_async_send_method = httpx .AsyncClient .send
108+ real_httpx_send_method = httpx .Client .send
108109
109110
110111# this method only returns a mock response if the URL accessed is weaviate.
111112# Otherwise, it passes the request through to real_httpx_send_method
112113# Important since Weaviate might make some PyPI requests (which don't need to be mocked).
113- async def mock_httpx_send_method (self , * args , ** kwargs ):
114- request = args [0 ]
115-
116- if "weaviate.cloud" not in str (request .url ):
117- return await real_httpx_send_method (self , * args , ** kwargs )
118-
114+ def _mock_weaviate_cloud_response (request ):
119115 assert WCD_URL in str (request .url )
120116 assert ("x-cohere-api-key" , COHERE_APIKEY ) in request .headers .items ()
121117 assert WCD_APIKEY in request .headers ["authorization" ]
@@ -127,7 +123,7 @@ async def mock_httpx_send_method(self, *args, **kwargs):
127123 # -H "x-weaviate-api-key: $WCD_API_KEY" \
128124 # -H "x-weaviate-cluster-url: $WCD_URL" \
129125 # -H "authorization: Bearer $WCD_API_KEY"
130- result = httpx .Response (
126+ return httpx .Response (
131127 status_code = 200 ,
132128 text = json .dumps (
133129 {
@@ -156,16 +152,40 @@ async def mock_httpx_send_method(self, *args, **kwargs):
156152 ),
157153 )
158154
159- return result
155+
156+ async def mock_httpx_async_send_method (self , * args , ** kwargs ):
157+ request = args [0 ]
158+
159+ if "weaviate.cloud" not in str (request .url ):
160+ return await real_httpx_async_send_method (self , * args , ** kwargs )
161+
162+ return _mock_weaviate_cloud_response (request )
163+
164+
165+ def mock_httpx_send_method (self , * args , ** kwargs ):
166+ request = kwargs .get ("request" ) or args [0 ]
167+
168+ if "weaviate.cloud" not in str (request .url ):
169+ return real_httpx_send_method (self , * args , ** kwargs )
170+
171+ return _mock_weaviate_cloud_response (request )
160172
161173
162174def test_cloud_resource ():
163175 """Mocks the httpx.AsyncClient.send method.
164176 Makes sure the url and headers (incl' authorization) of the request made by Weaviate are the same as
165177 specified in the Resource config.
166178 """
167- with mock .patch .object (
168- httpx .AsyncClient , "send" , autospec = True , side_effect = mock_httpx_send_method
179+ with (
180+ mock .patch .object (
181+ httpx .AsyncClient ,
182+ "send" ,
183+ autospec = True ,
184+ side_effect = mock_httpx_async_send_method ,
185+ ),
186+ mock .patch .object (
187+ httpx .Client , "send" , autospec = True , side_effect = mock_httpx_send_method
188+ ),
169189 ):
170190
171191 @asset
0 commit comments