unki/integration/docker_compose.py

import subprocess
import shlex
from utils import *
from docker import docker, check_state, print_errors


class Compose:
    ALL_CONTAINERS = [
        'u_agent',
        'u_server',
        'u_db',
        'tests_runner',
    ]

    def __init__(self):
        self.container_tpl = 'integration_%s_%d'
        self.cmd_container = self.container_tpl % ('tests_runner', 1)
        self.ALL_CONTAINERS = [self.container_tpl % (c, 1) for c in self.ALL_CONTAINERS]
        self.scaled_svc = {}
        self.scale("u_agent", 2)

    def scale(self, svc, count):
        for c in range(1, count):
            new_container = self.container_tpl % (svc, c + 1)
            self.ALL_CONTAINERS.append(new_container)
        self.scaled_svc[svc] = count

    def _call(self, *args):
        cmd = [
            'docker-compose',
            '--no-ansi',
        ] + list(args)
        log(f'Running docker-compose command: {cmd}')
        subprocess.check_call(cmd)

    def up(self):
        log('Instanciating cluster')
        scaled = [f"{k}={v}" for k, v in self.scaled_svc.items()]
        if len(scaled) > 0:
            scaled.insert(0, '--scale')
        self._call('up', '-d', *scaled)

    def down(self):
        log('Shutting down cluster')
        self._call('down')

    def stop(self):
        log('Stopping cluster')
        self._call('stop')

    def run(self, cmd):
        container = self.cmd_container
        if isinstance(cmd, str):
            cmd = shlex.split(cmd)
        result = docker([
            'exec',
            '-ti',
            container
        ] + cmd)
        return result

    def is_alive(self):
        log('Check if all containers are alive')

        errors = check_state(self.ALL_CONTAINERS)

        if errors:
            print_errors(errors)
            raise TestsError('Error during `is_alive` check')
        else:
            log('All containers are alive')