fix: fix tests

jina-ai · Jul 23, 2024 · 829a6d2 · 829a6d2
1 parent 42d9c35
commit 829a6d2
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/jina/clients/base/grpc.py b/jina/clients/base/grpc.py
@@ -145,7 +145,7 @@ async def _get_results(
                                 async for (
                                     response
                                 ) in stream_rpc.stream_rpc_with_retry():
-                                    yield response
+                                    yield response, None
                             else:
                                 unary_rpc = UnaryRpc(
                                     channel=channel,
@@ -169,7 +169,7 @@ async def _get_results(
                                     **kwargs,
                                 )
                                 async for response in unary_rpc.unary_rpc_with_retry():
-                                    yield response
+                                    yield response, None
                         except (grpc.aio.AioRpcError, InternalNetworkError) as err:
                             await self._handle_error_and_metadata(err)
         except KeyboardInterrupt:

diff --git a/jina/clients/base/websocket.py b/jina/clients/base/websocket.py
@@ -217,7 +217,7 @@ def _request_handler(
                     )
                     if self.show_progress:
                         p_bar.update()
-                    yield response
+                    yield response, None
             except Exception as ex:
                 exception_raised = ex
                 try:

diff --git a/jina/clients/mixin.py b/jina/clients/mixin.py
@@ -412,10 +412,11 @@ async def _get_results(*args, **kwargs):
                 if return_results:
                     resp.document_array_cls = inferred_return_type
                     if return_responses:
-                        resp.data.docs = da
+                        if da is not None:
+                            resp.data.docs = da
                         result.append(resp)
                     else:
-                        result.extend(da)
+                        result.extend(da if da is not None else resp.data.docs)
 
             if return_results:
                 if not return_responses and is_singleton and len(result) == 1:
@@ -540,13 +541,14 @@ async def post(
                     is_singleton = True
                     result.document_array_cls = DocList[return_type]
             if not return_responses:
-                ret_docs = da
+                ret_docs = da if da is not None else result.data.docs
                 if is_singleton and len(ret_docs) == 1:
                     yield ret_docs[0]
                 else:
                     yield ret_docs
             else:
-                result.data.docs = da
+                if da is not None:
+                    result.data.docs = da
                 yield result
 
     async def stream_doc(