service.py

import dataclasses
import os
import pathlib
 
from testsuite.environment import service
from testsuite.environment import utils
 
DEFAULT_HOST = 'localhost'
DEFAULT_GRPC_TLS_PORT = 2135
DEFAULT_GRPC_PORT = 2136
DEFAULT_MON_PORT = 8765
DEFAULT_DATABASE = 'local'
DEFAULT_CONTAINER_NAME = 'ydb-local-testsuite'
# TODO (lemito): change to ydbplatform/local-ydb:latest
DEFAULT_DOCKER_IMAGE = 'cr.yandex/yc/yandex-docker-local-ydb:latest'
 
PLUGIN_DIR = pathlib.Path(__file__).parent
SCRIPTS_DIR = PLUGIN_DIR.joinpath('scripts')
 
 
@dataclasses.dataclass(frozen=True)
class ServiceSettings:
    host: str
    grpc_port: int
    mon_port: int
    ic_port: int
    database: str
    wait_time: int = 30
 
 
def create_ydb_service(
    service_name: str,
    working_dir: str,
    settings: ServiceSettings | None = None,
    env: dict[str, str] | None = None,
):
    if settings is None:
        settings = get_service_settings()
    return service.ScriptService(
        service_name=service_name,
        script_path=str(SCRIPTS_DIR.joinpath('service-ydb.sh')),
        working_dir=working_dir,
        environment={
            'YDB_TMPDIR': working_dir,
            'YDB_SCRIPTS_DIR': str(SCRIPTS_DIR),
            'YDB_HOSTNAME': settings.host,
            'YDB_GRPC_PORT': str(settings.grpc_port),
            'YDB_GRPC_TLS_PORT': str(DEFAULT_GRPC_TLS_PORT),
            'YDB_MON_PORT': str(settings.mon_port),
            'YDB_CONTAINER_NAME': DEFAULT_CONTAINER_NAME,
            'YDB_DOCKER_IMAGE': DEFAULT_DOCKER_IMAGE,
            'YDB_WAIT_TIME': str(settings.wait_time),
            **(env or {}),
        },
        check_ports=[settings.grpc_port, settings.mon_port],
    )
 
 
def get_service_settings():
    return ServiceSettings(
        host=os.getenv('TESTSUITE_YDB_HOST', DEFAULT_HOST),
        grpc_port=utils.getenv_int(
            'TESTSUITE_YDB_GRPC_PORT',
            DEFAULT_GRPC_PORT,
        ),
        mon_port=utils.getenv_int('TESTSUITE_YDB_MON_PORT', DEFAULT_MON_PORT),
        ic_port=utils.getenv_int('TESTSUITE_YDB_IC_PORT', 0),
        database=os.getenv('TESTSUITE_YDB_DATABASE', DEFAULT_DATABASE),
        wait_time=utils.getenv_int('TESTSUITE_YDB_WAIT_TIME', 30),
    )