class MagicChecker(EventEmitter):
"Test magic against a single ip:port."
connect_timeout = 3
wait_timeout = 5
def __init__(self, magic):
"""
Given an IP address and port, try the magic on it, emitting a 'result' event.
"""
EventEmitter.__init__(self)
self.magic = magic + "\n"
self.timeouts = []
self.conn = None
self.output = ''
self.start = 0
self.latency = 0
self.client = TcpClient()
self.client.on("connect", self.connect)
self.client.on("connect_error", self.connect_error)
def check(self, host, port):
self.host = host
self.port = port
self.conn_start = time()
self.client.connect(host, port, connect_timeout=self.connect_timeout)
def connect(self, conn):
self.conn = conn
self.start = time()
self.latency = self.start - self.conn_start
self.timeouts.append(
schedule(self.wait_timeout, self.report, "TIMEOUT", "wait")
)
conn.on("data", self.data)
conn.on("close", self.close)
conn.write(self.magic)
conn.pause(False)
def connect_error(self, err_type, err_id, err_str):
self.start = time()
self.latency = 0
self.report("CONN_ERR", err_str)
def data(self, chunk):
self.output += chunk
def close(self):
self.report("CLOSE", '')
def report(self, result, details):
if (result, details) == ("TIMEOUT", "connect"):
wait = time() - self.conn_start
else:
wait = time() - self.start - self.latency
for timeout in self.timeouts:
timeout.delete()
if self.conn:
self.conn.close()
self.emit("result",
result, details, self.output, wait, self.host, self.port)
class TestTcpClientConnect(unittest.TestCase):
def setUp(self):
self.loop = loop.make()
self.connect_count = 0
self.error_count = 0
self.last_error_type = None
self.last_error = None
self.timeout_hit = False
self.conn = None
def check_connect(conn):
self.conn = conn
self.assertTrue(conn.tcp_connected)
self.connect_count += 1
conn.write(b"test")
conn.close()
self.loop.schedule(1, self.loop.stop)
def check_error(err_type, err_id, err_str):
self.error_count += 1
self.last_error_type = err_type
self.last_error = err_id
self.loop.schedule(1, self.loop.stop)
def timeout():
self.loop.stop()
self.timeout_hit = True
self.timeout = timeout
self.client = TcpClient(self.loop)
self.client.on('connect', check_connect)
self.client.on('connect_error', check_error)
def test_connect(self):
self.server = LittleServer(
(test_host, test_port),
LittleRequestHandler
)
t = threading.Thread(target=self.server.serve_forever)
t.setDaemon(True)
t.start()
self.client.connect(test_host, test_port)
self.loop.schedule(2, self.timeout)
self.loop.run()
self.assertFalse(self.conn.tcp_connected)
self.assertEqual(self.connect_count, 1)
self.assertEqual(self.error_count, 0)
self.assertEqual(self.timeout_hit, False)
self.server.shutdown()
self.server.socket.close()
def test_connect_refused(self):
self.client.connect(test_host, test_port + 1)
self.loop.schedule(3, self.timeout)
self.loop.run()
self.assertEqual(self.connect_count, 0)
self.assertEqual(self.error_count, 1)
self.assertEqual(self.last_error_type, socket.error)
self.assertEqual(self.last_error, errno.ECONNREFUSED)
self.assertEqual(self.timeout_hit, False)
def test_connect_noname(self):
self.client.connect('does.not.exist', test_port)
self.loop.schedule(3, self.timeout)
self.loop.run()
self.assertEqual(self.connect_count, 0)
self.assertEqual(self.error_count, 1)
self.assertEqual(self.last_error_type, socket.gaierror)
self.assertEqual(self.last_error, socket.EAI_NONAME)
self.assertEqual(self.timeout_hit, False)
def test_connect_timeout(self):
self.client.connect('128.66.0.1', test_port, 1)
self.loop.schedule(3, self.timeout)
self.loop.run()
self.assertEqual(self.connect_count, 0)
self.assertEqual(self.error_count, 1)
self.assertEqual(self.last_error_type, socket.error)
self.assertEqual(self.last_error, errno.ETIMEDOUT,
errno.errorcode.get(self.last_error, self.last_error))
self.assertEqual(self.timeout_hit, False)