diff --git a/tests/testlib.py b/tests/testlib.py index d173c378..ace8f0a2 100644 --- a/tests/testlib.py +++ b/tests/testlib.py @@ -406,24 +406,6 @@ def get_docker_host(): class DockerizedSshDaemon(object): - mitogen_test_distro = os.environ.get('MITOGEN_TEST_DISTRO', 'debian') - if '-' in mitogen_test_distro: - distro, _py3 = mitogen_test_distro.split('-') - else: - distro = mitogen_test_distro - _py3 = None - - if _py3 == 'py3': - python_path = '/usr/bin/python3' - else: - python_path = '/usr/bin/python' - - image = 'mitogen/%s-test' % (distro,) - - # 22/tcp -> 0.0.0.0:32771 - PORT_RE = re.compile(r'([^/]+)/([^ ]+) -> ([^:]+):(.*)') - port = None - def _get_container_port(self): s = subprocess__check_output(['docker', 'port', self.container_name]) for line in s.decode().splitlines(): @@ -454,7 +436,24 @@ class DockerizedSshDaemon(object): subprocess__check_output(args) self._get_container_port() - def __init__(self): + def __init__(self, mitogen_test_distro=os.environ.get('MITOGEN_TEST_DISTRO', 'debian')): + if '-' in mitogen_test_distro: + distro, _py3 = mitogen_test_distro.split('-') + else: + distro = mitogen_test_distro + _py3 = None + + if _py3 == 'py3': + self.python_path = '/usr/bin/python3' + else: + self.python_path = '/usr/bin/python' + + self.image = 'mitogen/%s-test' % (distro,) + + # 22/tcp -> 0.0.0.0:32771 + self.PORT_RE = re.compile(r'([^/]+)/([^ ]+) -> ([^:]+):(.*)') + self.port = None + self.start_container() def get_host(self): @@ -521,7 +520,13 @@ class DockerMixin(RouterMixin): super(DockerMixin, cls).setUpClass() if os.environ.get('SKIP_DOCKER_TESTS'): raise unittest2.SkipTest('SKIP_DOCKER_TESTS is set') - cls.dockerized_ssh = DockerizedSshDaemon() + + # we want to be able to override test distro for some tests that need a different container spun up + daemon_args = {} + if hasattr(cls, 'mitogen_test_distro'): + daemon_args['mitogen_test_distro'] = cls.mitogen_test_distro + + cls.dockerized_ssh = DockerizedSshDaemon(**daemon_args) cls.dockerized_ssh.wait_for_sshd() @classmethod