unki/integration/integration_tests.py

import signal
import sys
from utils import *
from docker import rebuild_images_if_needed
from docker_compose import Compose


cluster = Compose()


def abort_handler(s, _):
    warn(f'Received signal: {s}')
    warn(f'Gracefully stopping...')
    cluster.down()


def usage_exit():
    usage = f"""Usage:
    python {__file__.split('/')[-1]} [--rebuild] [--preserve] [--no-run]"""

    print(usage)
    sys.exit(1)


def run_tests():
    allowed_args = set(["--rebuild", "--preserve", "--no-run"])
    args = sys.argv[1:]
    if not set(args).issubset(allowed_args):
        usage_exit()
    force_rebuild = '--rebuild' in args
    preserve_containers = '--preserve' in args
    only_setup_cluster = '--no-run' in args
    for s in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        signal.signal(s, abort_handler)
    rebuild_images_if_needed(force_rebuild)
    try:
        cluster.up()
        cluster.is_alive()
        if not only_setup_cluster:
            cluster.run('cargo test --test integration')
    except Exception as e:
        err(e)
        sys.exit(1)
    finally:
        if not preserve_containers and not only_setup_cluster:
            cluster.down()


if __name__ == '__main__':
    run_tests()