Control a background process during a test (typically for an integration test)

pytest_background_proc.py

import os
import subprocess
import tempfile
import time

import pytest


ENABLE = os.environ.get('WITH_INTEGRATION_TESTS')
integration_test = pytest.mark.skipif(not ENABLE, reason="Integration tests skipped")


class BackgroundProc:
    START_TIMEOUT = 10
    EXIT_TIMEOUT = 10

    def __init__(self, cmd, env=None, expect_exitcode=0):
        self._cmd = cmd
        self._env = os.environ if env is None else env
        self._outfile = None
        self._proc = None
        self._expect_exitcode = expect_exitcode

    def __enter__(self):
        self.launch()

    def __exit__(self, *args):
        self.teardown()

    def launch(self):
        print(f'Launching process {self._cmd}')

        self._outfile = tempfile.NamedTemporaryFile()

        try:
            self._proc = subprocess.Popen(self._cmd, env=self._env, stderr=subprocess.STDOUT, stdout=self._outfile)
        except FileNotFoundError:
            path = os.environ['PATH']
            print(f'Failed to find {self._cmd} in PATH={path}')
            raise

        if not self._expect_in_output('Application is running'):
            raise RuntimeError(f"Process {self._cmd} didn't start in time")

        returncode = self._proc.poll()
        if returncode is not None:
            raise RuntimeError(f'Process {self._cmd} crashed (code={returncode})')

        print(f'Process {self._cmd} started properly')

    def teardown(self):
        print(f'Killing process {self._cmd}')

        if self._proc.poll() is None:
            self._proc.terminate()
        else:
            self._print_output()
            raise RuntimeError(f'Process {self._cmd} crashed before teardown (code={returncode})')

        returncode = self._wait_for_exit()

        self._print_output()

        if returncode is None:
            raise RuntimeError(f"Process {self._cmd} didn't close in time")

        if returncode != self._expect_exitcode:
            raise RuntimeError(f'Process {self._cmd} exited (code={returncode}, expected={self._expect_exitcode})')

        print(f'Process {self._cmd} closed properly (code={returncode})')

    def _wait_for_exit(self):
        for _ in range(int(self.START_TIMEOUT / 0.1)):
            time.sleep(0.1)
            self._proc.poll()
            if self._proc.returncode is not None:
                return self._proc.returncode

    def _expect_in_output(self, substr):
        for _ in range(int(self.EXIT_TIMEOUT / 0.1)):
            time.sleep(0.1)
            with open(self._outfile.name) as fh:
                if substr in fh.read():
                    return True
        return False

    def _print_output(self):
        with open(self._outfile.name) as fh:
            for line in fh:
                print(f'{self._cmd}: {line}', end='')