fix: user session.request

jina-ai · Dec 5, 2023 · 7fbe42f · 7fbe42f
1 parent 42f444a
commit 7fbe42f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/jina/serve/runtimes/gateway/request_handling.py b/jina/serve/runtimes/gateway/request_handling.py
@@ -170,8 +170,8 @@ async def _load_balance(self, request):
                 except Exception:
                     self.logger.debug('No JSON payload found in request')
 
-                async with _RequestContextManager(
-                    session._request(request.method, target_url, **request_kwargs)
+                async with session.request(
+                    request.method, target_url, **request_kwargs
                 ) as response:
                     # Looking for application/octet-stream, text/event-stream, text/stream
                     if response.content_type.endswith('stream'):
@@ -197,6 +197,7 @@ async def _load_balance(self, request):
                         body=content,
                         status=response.status,
                         content_type=response.content_type,
+                        headers=response.headers,
                     )
         except aiohttp.ClientError as e:
             return web.Response(text=f'Error: {str(e)}', status=500)

diff --git a/tests/integration/docarray_v2/test_issues.py b/tests/integration/docarray_v2/test_issues.py
@@ -184,13 +184,11 @@ async def test_issue_6090_get_params(streaming_deployment):
 
         async with session.get(url) as resp:
             async for chunk in resp.content.iter_any():
-                print(chunk)
                 events = chunk.split(b'event: ')[1:]
                 for event in events:
                     if event.startswith(b'update'):
-                        parsed = event[HTTPClientlet.UPDATE_EVENT_PREFIX:].decode()
+                        parsed = event[HTTPClientlet.UPDATE_EVENT_PREFIX :].decode()
                         parsed = SimpleInput.parse_raw(parsed)
-                        print(parsed)
                         docs.append(parsed)
                     elif event.startswith(b'end'):
                         pass