unki/integration/docker.py

import subprocess
import shlex

from utils import *

BASE_IMAGE_DIR = '../images/integration-tests'

DOCKERFILES = [
    {
        'name': 'u_agent',
        'ctx': BASE_IMAGE_DIR,
    },
    {
        'name': 'u_server',
        'ctx': BASE_IMAGE_DIR,
    },
    {
        'name': 'u_db',
        'ctx': BASE_IMAGE_DIR,
    },
    {
        'name': 'tests_runner',
        'ctx': BASE_IMAGE_DIR,
    },
]


def docker(args):
    cmd = ['docker'] + args
    log(f'Running docker command: {cmd}')
    return subprocess.run(
        cmd,
        check=True,
    )


def print_errors(errors):
    err_msg = '\n'.join(
        '  {container}: {error}'.format(container=item['container'],
                                        error=item['error'])
        for item in errors)

    err('There are some errors in next containers:\n%s' % err_msg)


def check_state(containers):
    errors = []
    for container in containers:
        ret, out = subprocess.getstatusoutput(
            'docker inspect --format \'{{ .State.Running }}\' %s'
            % container)
        out = out.strip()
        if ret == 0:
            if out == 'true':
                continue
            else:
                errors.append({'container': container,
                               'error': 'Bad state: Running=%s' % out})
        else:
            errors.append({'container': container,
                           'error': out})

    return errors


def rebuild_images_if_needed(force_rebuild=False):
    for img in DOCKERFILES:
        ctx = img['ctx']
        name = img.get('name')
        df_suffix = 'Dockerfile'
        img_name = f'unki/{name}'
        log(f'Building docker image {img_name}')
        cmd = [
            'build',
            '-t', img_name,
            '-f', f'{BASE_IMAGE_DIR}/{name}.{df_suffix}',
            ctx,
        ]
        if force_rebuild:
            cmd += ['--no-cache']
        docker(cmd)


class Compose:
    ALL_IMAGES = [
        '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_IMAGES]
        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',
            '--ansi=never',
        ] + list(args)
        log(f'Running docker-compose command: {cmd}')
        subprocess.check_call(cmd)

    def up(self):
        log(f'Instanciating cluster: {self.ALL_CONTAINERS}')
        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')

    def print_containers_logs(self):
        for container in self.ALL_CONTAINERS:
            try:
                docker([
                    'logs',
                    container
                ])
            except Exception:
                pass