import os import random import re import socket import sys import time import unittest import urlparse import mitogen.master if mitogen.is_master: # TODO: shouldn't be necessary. import docker DATA_DIR = os.path.join(os.path.dirname(__file__), 'data') sys.path.append(DATA_DIR) def set_debug(): import logging logging.getLogger('mitogen').setLevel(logging.DEBUG) def data_path(suffix): path = os.path.join(DATA_DIR, suffix) if path.endswith('.key'): # SSH is funny about private key permissions. os.chmod(path, 0600) return path def wait_for_port( host, port, pattern=None, connect_timeout=0.5, receive_timeout=0.5, overall_timeout=5.0, sleep=0.1, ): """Attempt to connect to host/port, for upto overall_timeout seconds. If a regex pattern is supplied try to find it in the initial data. Return None on success, or raise on error. """ start = time.time() end = start + overall_timeout addr = (host, port) while time.time() < end: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(connect_timeout) try: sock.connect(addr) except socket.error: # Failed to connect. So wait then retry. time.sleep(sleep) continue if not pattern: # Success: We connected & there's no banner check to perform. sock.shutdown(socket.SHUTD_RDWR) sock.close() return sock.settimeout(receive_timeout) data = '' found = False while time.time() < end: try: resp = sock.recv(1024) except socket.timeout: # Server stayed up, but had no data. Retry the recv(). continue if not resp: # Server went away. Wait then retry the connection. time.sleep(sleep) break data += resp if re.search(pattern, data): found = True break sock.shutdown(socket.SHUT_RDWR) sock.close() if found: # Success: We received the banner & found the desired pattern return else: # Failure: The overall timeout expired if pattern: raise socket.timeout('Timed out while searching for %r from %s:%s' % (pattern, host, port)) else: raise socket.timeout('Timed out while connecting to %s:%s' % (host, port)) class TestCase(unittest.TestCase): def assertRaises(self, exc, func, *args, **kwargs): """Like regular assertRaises, except return the exception that was raised. Can't use context manager because tests must run on Python2.4""" try: func(*args, **kwargs) except exc, e: return e except BaseException, e: assert 0, '%r raised %r, not %r' % (func, e, exc) assert 0, '%r did not raise %r' % (func, exc) class DockerizedSshDaemon(object): def __init__(self): self.docker = docker.from_env(version='auto') self.container_name = 'mitogen-test-%08x' % (random.getrandbits(64),) self.container = self.docker.containers.run( image='d2mw/mitogen-test', detach=True, publish_all_ports=True, ) self.container.reload() self.port = (self.container.attrs['NetworkSettings']['Ports'] ['22/tcp'][0]['HostPort']) self.host = self.get_host() def get_host(self): if self.docker.api.base_url == 'http+docker://localunixsocket': return 'localhost' parsed = urlparse.urlparse(self.docker.api.base_url) return parsed.netloc.partition(':')[0] def wait_for_sshd(self): wait_for_port(self.get_host(), int(self.port), pattern='OpenSSH') def close(self): self.container.stop() self.container.remove() class RouterMixin(object): broker_class = mitogen.master.Broker router_class = mitogen.master.Router def setUp(self): super(RouterMixin, self).setUp() self.broker = self.broker_class() self.router = self.router_class(self.broker) def tearDown(self): self.broker.shutdown() self.broker.join() super(RouterMixin, self).tearDown() class DockerMixin(RouterMixin): @classmethod def setUpClass(cls): super(DockerMixin, cls).setUpClass() cls.dockerized_ssh = DockerizedSshDaemon() cls.dockerized_ssh.wait_for_sshd() @classmethod def tearDownClass(cls): cls.dockerized_ssh.close() super(DockerMixin, cls).tearDownClass() def docker_ssh(self, **kwargs): kwargs.setdefault('hostname', self.dockerized_ssh.host) kwargs.setdefault('port', self.dockerized_ssh.port) kwargs.setdefault('check_host_keys', False) return self.router.ssh(**kwargs) def docker_ssh_any(self, **kwargs): return self.docker_ssh( username='has-sudo-nopw', password='y', )