Repository: jina-ai/jina
File: tests/unit/serve/gateway/test_gateway.py (Python)

Code before:
            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
        ),
        (
            _dummy_fastapi_gateway_yaml_path,
            {'arg1': 'arg1'},
            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
        ),
    ],
)
def test_custom_gateway_no_executors(uses, uses_with, expected):
    worker_port, worker_process = _start_worker_runtime('ProcessExecutor')
    gateway_port, gateway_process = _start_gateway_runtime(uses, uses_with, worker_port)
    _validate_dummy_custom_gateway_response(gateway_port, expected)
    _validate_custom_gateway_process(
        gateway_port, 'hello', {'text': 'helloworld', 'tags': {'processed': True}}
    )
    gateway_process.terminate()
    gateway_process.join()
    worker_process.terminate()
    worker_process.join()

    assert gateway_process.exitcode == 0
    assert worker_process.exitcode == 0


def test_stream_individual_executor():
    from docarray import DocumentArray, Document

    from jina.serve.runtimes.gateway.http.fastapi import FastAPIBaseGateway
    from jina import Flow, Executor, requests


    PARAMETERS = {'dog': 'woof'}

    class MyGateway(FastAPIBaseGateway):
        @property
        def app(self):
            from fastapi import FastAPI

            app = FastAPI(title='Custom FastAPI Gateway')

            @app.get("/endpoint")
            async def get(text: str):
                docs = await self.executor['executor1'].post(on='/', inputs=DocumentArray([Document(text=text), Document(text=text.upper())]), parameters=PARAMETERS)
                return {'result': [doc.text for doc in docs]}

            return app

    class FirstExec(Executor):
        @requests
        def func(self, docs, **kwargs):

Review comment (suggestion):
let's add parameters `n_shards` and `n_replicas` in order to test how this feature supports replicas and shards

Code after:
            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
        ),
        (
            _dummy_fastapi_gateway_yaml_path,
            {'arg1': 'arg1'},
            {'arg1': 'arg1', 'arg2': 'world', 'arg3': 'default-arg3'},
        ),
    ],
)
def test_custom_gateway_no_executors(uses, uses_with, expected):
    worker_port, worker_process = _start_worker_runtime('ProcessExecutor')
    gateway_port, gateway_process = _start_gateway_runtime(uses, uses_with, worker_port)
    _validate_dummy_custom_gateway_response(gateway_port, expected)
    _validate_custom_gateway_process(
        gateway_port, 'hello', {'text': 'helloworld', 'tags': {'processed': True}}
    )
    gateway_process.terminate()
    gateway_process.join()
    worker_process.terminate()
    worker_process.join()

    assert gateway_process.exitcode == 0
    assert worker_process.exitcode == 0


def test_stream_individual_executor_simple():
    from docarray import DocumentArray, Document

    from jina.serve.runtimes.gateway.http.fastapi import FastAPIBaseGateway
    from jina import Flow, Executor, requests


    PARAMETERS = {'dog': 'woof'}

    class MyGateway(FastAPIBaseGateway):
        @property
        def app(self):
            from fastapi import FastAPI

            app = FastAPI(title='Custom FastAPI Gateway')

            @app.get('/endpoint')
            async def get(text: str):
                docs = await self.executor['executor1'].post(on='/', inputs=DocumentArray([Document(text=text), Document(text=text.upper())]), parameters=PARAMETERS)
                return {'result': docs.texts}

            return app

    class FirstExec(Executor):
        @requests
        def func(self, docs, **kwargs):