from jina import Flow, Executor, requests
from docarray import BaseDocument, DocumentArray
from docarray.documents import Image
from docarray.typing import AnyTensor

class InputDoc(BaseDocument):
    img: Image

class OutputDoc(BaseDocument):
    embedding: AnyTensor

class MyExec(Executor):
    @requests(on='/bar')
    def bar(
        self, docs: DocumentArray[InputDoc], **kwargs
    ) -> DocumentArray[OutputDoc]:

        docs_return = DocumentArray[OutputDoc](
            [OutputDoc(embedding=np.zeros((100, 1))) for _ in range(len(docs))]
        )

        return docs_return

with Flow().add(uses=MyExec) as f:
    docs = f.post(
        on='/bar',
        inputs=InputDoc(img=Image(tensor=np.zeros((3, 224, 224)))),
        return_type=DocumentArray[OutputDoc],
    )
    assert docs[0].embedding.shape == (100, 1)
    assert docs.__class__.document_type == OutputDoc

import pytest

from fastapi import FastAPI
from httpx import AsyncClient

from docarray import BaseDocument
from docarray.base_document import DocumentResponse
from docarray.documents import Image
from docarray.typing import AnyTensor


import numpy as np


class InputDoc(BaseDocument):
    img: Image


class OutputDoc(BaseDocument):
    embedding: AnyTensor


input_doc = InputDoc(img=Image(tensor=np.zeros((3, 224, 224))))

app = FastAPI()


@app.post("/embed/", response_model=OutputDoc, response_class=DocumentResponse)
async def embed(doc: InputDoc) -> OutputDoc:
    doc = OutputDoc(embedding=np.zeros((100, 1)))
    return doc


async with AsyncClient(app=app, base_url="http://test") as ac:
    response = await ac.post("/embed/", data=input_doc.json())

doc = OutputDoc.parse_raw(response.content.decode())
assert doc.embedding.shape == (100, 1)