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', '--ansi=never', ] + 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')