You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
mitogen/tests/testlib.py

185 lines
5.2 KiB
Python

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)
def assertContains(self, needle, hay):
assert needle in hay, "%r not found in %r" % (needle, hay)
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',
)