Test harness for a Vault server

vault_server.py

import logging
import subprocess
import sys
import time


log = logging.getLogger(__name__)


class VaultServer(object):
    def __init__(self,
                 path='vault',
                 log_level='debug',
                 log_format='',
                 config='',
                 dev_root_token_id='token',
                 dev_listen_address='127.0.0.1:8200',
                 dev_plugin_dir='',
                 dev_no_store_token=False,
                 startup_time_sec=1
                 ):
        cmd = [
            path,
            'server',
            '-dev',
            '-log-level', log_level,
            '-dev-root-token-id', dev_root_token_id,
            '-dev-listen-address', dev_listen_address
        ]
        if log_format:
            cmd.extend(['-log-format', log_format])
        if dev_plugin_dir:
            cmd.extend(['-dev-plugin-dir', dev_plugin_dir])
        if config:
            cmd.extend(['-config', config])
        if dev_no_store_token:
            cmd.extend(['-dev-no-store-token'])
        self.process = subprocess.Popen(cmd)
        self.vault_addr = f'http://{dev_listen_address}'
        time.sleep(startup_time_sec)


    def __enter__(self):
        self.process.poll()
        if self.process.returncode is not None:
            log.fatal(f'Vault has exited with code {self.process.returncode}')
            sys.exit(self.process.returncode)
        return self.vault_addr

    def __exit__(self, type, value, traceback):
        self.process.terminate()
        while self.process.returncode is None:
            self.process.poll()