test: add test with shaped ndarray (jina-ai#5900)

Co-authored-by: Jina Dev Bot <[email protected]>

Author: JoanFM

jina_cli/autocomplete.py

@@ -247,7 +247,6 @@
         'cloud scale': ['--help', '--executor', '--replicas'],
         'cloud recreate': ['--help'],
         'cloud logs': ['--help', '--gateway', '--executor'],
-        'cloud survey': ['--help'],
         'cloud': [
             '--help',
             '--version',
@@ -267,7 +266,6 @@
             'scale',
             'recreate',
             'logs',
-            'survey',
         ],
         'help': ['--help'],
         'pod': [

tests/integration/docarray_v2/test_v2.py

@@ -5,7 +5,7 @@
 import numpy as np
 from docarray import BaseDoc, DocList
 from docarray.documents import ImageDoc
-from docarray.typing import AnyTensor, ImageUrl
+from docarray.typing import AnyTensor, ImageUrl, NdArray
 from docarray.documents import TextDoc
 from docarray.documents.legacy import LegacyDocument
 from jina.helper import random_port
@@ -292,7 +292,6 @@ def foo(self, docs: DocList[ProcessingTestDocConditions], **kwargs) -> DocList[P
                     fp.write(doc.text)
                     doc.text += f' processed by {self.metas.name}'
 
-
     class FirstExec(Executor):
         @requests
         def foo(self, docs: DocList[LegacyDocument], **kwargs) -> DocList[ProcessingTestDocConditions]:
@@ -462,6 +461,32 @@ def foo(self, docs: DocList[TextDoc], **kwargs) -> DocList[TextDoc]:
         assert len(ret) == 0
 
 
+@pytest.mark.parametrize('protocol', ['grpc', 'http', 'websocket'])
+@pytest.mark.parametrize('ctxt_manager', ['deployment', 'flow'])
+def test_input_output_with_shaped_tensor(protocol, ctxt_manager):
+    if ctxt_manager == 'deployment' and protocol == 'websocket':
+        return
+
+    class MyDoc(BaseDoc):
+        text: str
+        embedding: NdArray[128]
+
+    class Foo(Executor):
+        @requests(on='/hello')
+        def foo(self, docs: DocList[MyDoc], **kwargs) -> DocList[MyDoc]:
+            for doc in docs:
+                doc.text += 'Processed by foo'
+
+    if ctxt_manager == 'flow':
+        ctxt_mgr = Flow(protocol=protocol).add(uses=Foo)
+    else:
+        ctxt_mgr = Deployment(protocol=protocol, uses=Foo)
+
+    with ctxt_mgr:
+        ret = ctxt_mgr.post(on='/hello', inputs=DocList[MyDoc]([MyDoc(text='', embedding=np.random.rand(128))]))
+        assert len(ret) == 1
+
+
 @pytest.mark.parametrize('protocol', ['grpc', 'http', 'websocket'])
 @pytest.mark.parametrize('ctxt_manager', ['deployment', 'flow'])
 def test_send_parameters(protocol, ctxt_manager):
@@ -651,14 +676,14 @@ class Previous(Executor):
         def foo(self, docs: DocList[TextDoc], **kwargs) -> DocList[TextDoc]:
             pass
 
-    f = Flow(protocol=protocol).add(uses=Previous, name='previous').add(uses=First, name='first', needs='previous').add(uses=Second, name='second', needs='previous').needs_all()
+    f = Flow(protocol=protocol).add(uses=Previous, name='previous').add(uses=First, name='first', needs='previous').add(
+        uses=Second, name='second', needs='previous').needs_all()
 
     with pytest.raises(RuntimeFailToStart):
         with f:
             pass
 
 
-
 class ExternalDeploymentDoc(BaseDoc):
     tags: Dict[str, str] = {}