Repository: jina-ai/jina
File: tests/unit/serve/executors/test_executor.py (Python)

Code before:
        assert set(os.listdir(tmpdir)) == {
            'executor0',
            'gateway',
        }

        with open(os.path.join(tmpdir, 'gateway', 'gateway.yml')) as f:
            gatewayyaml = list(yaml.safe_load_all(f))[-1]
            assert (
                gatewayyaml['spec']['template']['spec']['containers'][0]['ports'][0][
                    'containerPort'
                ]
                == 2020
            )
            gateway_args = gatewayyaml['spec']['template']['spec']['containers'][0][
                'args'
            ]
            assert gateway_args[gateway_args.index('--port') + 1] == '2020'


@pytest.mark.parametrize(
    'exec_type',
    [Executor.StandaloneExecutorType.EXTERNAL, Executor.StandaloneExecutorType.SHARED],
)
@pytest.mark.parametrize(
    'uses',
    ['jinahub+docker://DummyHubExecutor', 'jinaai+docker://jina-ai/DummyHubExecutor'],
)
def test_to_docker_compose_yaml(tmpdir, exec_type, uses):
    compose_file = os.path.join(tmpdir, 'compose.yml')
    Executor.to_docker_compose_yaml(
        output_path=compose_file,
        port_expose=2020,
        uses=uses,
        executor_type=exec_type,
    )

    with open(compose_file) as f:
        services = list(yaml.safe_load_all(f))[0]['services']
        assert services['executor0']['image'].startswith('jinahub/')

        if exec_type == Executor.StandaloneExecutorType.SHARED:
            assert len(services) == 1
        else:
            assert len(services) == 2
            assert services['gateway']['ports'][0] == '2020:2020'
            gateway_args = services['gateway']['command']
            assert gateway_args[gateway_args.index('--port') + 1] == '2020'


def _create_test_data_message(counter=0):
    return list(request_generator('/', DocumentArray([Document(text=str(counter))])))[0]

Review comment (suggestion):
```suggestion
    ['jinaai+docker://jina-ai/DummyHubExecutor'],
```

Code after:
        assert set(os.listdir(tmpdir)) == {
            'executor0',
            'gateway',
        }

        with open(os.path.join(tmpdir, 'gateway', 'gateway.yml')) as f:
            gatewayyaml = list(yaml.safe_load_all(f))[-1]
            assert (
                gatewayyaml['spec']['template']['spec']['containers'][0]['ports'][0][
                    'containerPort'
                ]
                == 2020
            )
            gateway_args = gatewayyaml['spec']['template']['spec']['containers'][0][
                'args'
            ]
            assert gateway_args[gateway_args.index('--port') + 1] == '2020'


@pytest.mark.parametrize(
    'exec_type',
    [Executor.StandaloneExecutorType.EXTERNAL, Executor.StandaloneExecutorType.SHARED],
)
@pytest.mark.parametrize(
    'uses',
    ['jinaai+docker://jina-ai/DummyHubExecutor'],
)
def test_to_docker_compose_yaml(tmpdir, exec_type, uses):
    compose_file = os.path.join(tmpdir, 'compose.yml')
    Executor.to_docker_compose_yaml(
        output_path=compose_file,
        port_expose=2020,
        uses=uses,
        executor_type=exec_type,
    )

    with open(compose_file) as f:
        services = list(yaml.safe_load_all(f))[0]['services']
        assert services['executor0']['image'].startswith('jinahub/')

        if exec_type == Executor.StandaloneExecutorType.SHARED:
            assert len(services) == 1
        else:
            assert len(services) == 2
            assert services['gateway']['ports'][0] == '2020:2020'
            gateway_args = services['gateway']['command']
            assert gateway_args[gateway_args.index('--port') + 1] == '2020'


def _create_test_data_message(counter=0):
    return list(request_generator('/', DocumentArray([Document(text=str(counter))])))[0]