add test

Author: blampe

Makefile

@@ -1,7 +1,7 @@
 project := tchannel
 
 flake8 := flake8
-pytest := py.test -s --tb short --cov-config .coveragerc --cov \
+pytest := py.test -sv --tb short --cov-config .coveragerc --cov \
           $(project) tests
 
 html_report := --cov-report html

requirements-test.txt

@@ -1,6 +1,7 @@
 # Testing and coverage
 pytest
 pytest-cov
+pytest-tornado
 
 # Test all the pythons
 tox

tchannel/tchannel.py

@@ -8,10 +8,10 @@
 import tornado.ioloop
 import tornado.iostream
 
-from .tornado.connection import TornadoConnection
 from .exceptions import InvalidMessageException
 from .messages import CallRequestMessage
-from .timeout import timeout
+from .tornado.connection import TornadoConnection
+from .tornado.timeout import timeout
 
 
 log = logging.getLogger('tchannel')
@@ -31,15 +31,15 @@ def __init__(self, process_name=None):
     @tornado.gen.coroutine
     def add_peer(self, hostport):
         if hostport in self.peers:
-            return
+            raise tornado.gen.Return(self.peers[hostport])
 
         peer = yield self.make_out_connection(hostport)
         self.peers[hostport] = peer
         raise tornado.gen.Return(peer)
 
     def remove_peer(self, hostport):
         # TODO: Connection cleanup
-        self.peers.pop(hostport)
+        return self.peers.pop(hostport)
 
     @tornado.gen.coroutine
     def get_peer(self, hostport):
@@ -50,10 +50,13 @@ def get_peer(self, hostport):
         raise tornado.gen.Return(peer)
 
     @tornado.gen.coroutine
-    def make_out_connection(self, hostport):
+    def make_out_connection(self, hostport, sock=None):
         host, port = hostport.rsplit(":", 1)
 
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+        # TODO: change this to tornado.tcpclient.TCPClient to do async DNS
+        # lookups.
         stream = tornado.iostream.IOStream(sock)
 
         log.debug("connecting to hostport %s", hostport)

tchannel/tcurl.py

@@ -40,6 +40,8 @@ def tcurl(host, endpoint, headers, body):
     print " arg2:", getattr(response, 'arg_2', None)
     print " arg3:", getattr(response, 'arg_3', None)
 
+    raise tornado.gen.Return(response)
+
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -104,5 +106,5 @@ def main():
     )
 
 
-if __name__ == '__main__':
+if __name__ == '__main__':  # pragma: no cover
     main()

tchannel/timeout.py

@@ -26,9 +26,6 @@ def write(self, data, callback=None):
         """
         return self._stream.write(bytes(data), callback=callback)
 
-    def close(self):
-        self._stream.close()
-
 
 class TornadoConnection(Connection):
     """Handle speaking TChannel over a Tornado connection."""
@@ -51,3 +48,6 @@ def await(self, callback=None):
 
     def handle_calls(self, handler):
         return self.await(callback=self.wrap(handler))
+
+    def close(self):
+        return self._connection._stream.close()

tchannel/tornado/connection.py

@@ -0,0 +1,17 @@
+from __future__ import absolute_import
+
+import contextlib2
+import tornado.ioloop
+
+
+@contextlib2.contextmanager
+def timeout(connection, seconds=2):
+
+    handle = tornado.ioloop.IOLoop.instance().call_later(
+        seconds,
+        connection.close,
+    )
+
+    yield handle
+
+    tornado.ioloop.IOLoop.instance().remove_timeout(handle)

tchannel/tornado/timeout.py

@@ -0,0 +1,24 @@
+from __future__ import absolute_import
+
+import socket
+
+import pytest
+
+from tests.integration.server_manager import ServerManager
+
+
+@pytest.yield_fixture
+def server_manager(random_open_port):
+    with ServerManager(random_open_port) as manager:
+        yield manager
+
+
+@pytest.fixture
+def random_open_port():
+    """Find and return a random open TCP port."""
+    sock = socket.socket(socket.AF_INET)
+    try:
+        sock.bind(('', 0))
+        return sock.getsockname()[1]
+    finally:
+        sock.close()

tests/integration/__init__.py

@@ -0,0 +1,132 @@
+from __future__ import absolute_import
+
+import socket
+import threading
+try:
+    import SocketServer
+except ImportError:
+    import socketserver as SocketServer
+from contextlib import contextmanager
+
+import tchannel.socket as tchannel
+import tchannel.messages as tmessage
+
+
+class Expectation(object):
+    """Represents an expectation for the ServerManager."""
+    def __init__(self, matcher):
+        assert matcher is not None
+
+        # expectation.matches(req) accepts a Message and returns True or
+        # False.
+        self.matches = matcher
+
+        # expectation.respond(context, connection) accepts the context and the
+        # connection and writes output to the connection.
+        self._respond = None
+
+    @classmethod
+    def messageType(cls, msg_typ):
+        """Build an expectation that expects a mesasge with the given type."""
+        return cls(lambda msg: msg.message_type == msg_typ)
+
+    @property
+    def respond(self):
+        # Do nothing if an action setter wasn't called.
+        if self._respond:
+            return self._respond
+        else:
+            return (lambda ctx, conn: None)
+
+    def and_return(self, resp):
+        """Write the given Message as a response."""
+        def respond(ctx, conn):
+            return conn.frame_and_write(resp)
+        self._respond = respond
+
+
+class ServerManager(object):
+    """Provides a dynamically configurable TChannel server."""
+    TIMEOUT = 0.15
+
+    def __init__(self, port, timeout=None):
+        manager = self
+        self.port = port
+        self.timeout = timeout or self.TIMEOUT
+
+        class Handler(SocketServer.BaseRequestHandler):
+            def setup(self):
+                self.request.settimeout(manager.timeout)
+                self.tchan_conn = tchannel.SocketConnection(self.request)
+
+            def handle(self):
+                (host, port) = self.request.getsockname()
+                self.tchan_conn.await_handshake(headers={
+                    'host_port': '%s:%s' % (host, port),
+                    'process_name': 'tchannel_server-%s' % port
+                })
+                self.tchan_conn.handle_calls(manager.handle_call)
+
+        self.server = SocketServer.TCPServer(("", port), Handler)
+        self.thread = None
+        self._expectations = []
+
+    def expect_ping(self):
+        """Expect a Ping request.
+
+        Returns an Expectation to allow setting the response behavior."""
+        exp = Expectation.messageType(
+            tmessage.PingRequestMessage.message_type
+        )
+
+        self._expectations.append(exp)
+        return exp
+
+    def expect_call_request(self, endpoint):
+        def matcher(message):
+            return (
+                message.message_type == tmessage.CallRequestMessage.message_type and
+                message.arg_1 == endpoint
+            )
+        exp = Expectation(matcher)
+        self._expectations.append(exp)
+        return exp
+
+    def handle_call(self, context, connection):
+        for exp in self._expectations:
+            if exp.matches(context.message):
+                exp.respond(context, connection)
+
+    @contextmanager
+    def client_connection(self):
+        """Get an initiated Connection to this TChannel server."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(self.timeout)
+        try:
+            (host, port) = ('localhost', self.port)
+            sock.connect((host, port))
+
+            conn = tchannel.SocketConnection(sock)
+            conn.initiate_handshake(headers={
+                'host_port': '%s:%s' % (host, port),
+                'process_name': 'tchannel_client-%s' % port
+            })
+            conn.await_handshake_reply()
+            yield conn
+        finally:
+            sock.close()
+
+    def start(self):
+        assert self.thread is None, 'server already started'
+        self.thread = threading.Thread(target=self.server.serve_forever)
+        self.thread.start()
+
+    def stop(self):
+        self.server.shutdown()
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, *args):
+        self.stop()

tests/integration/conftest.py

@@ -1,151 +1,67 @@
 from __future__ import absolute_import
 
-import socket
-import threading
-try:
-    import SocketServer
-except ImportError:
-    import socketserver as SocketServer
-from contextlib import contextmanager
-
 import pytest
 
-import tchannel.socket as tchannel
 import tchannel.messages as tmessage
+from tchannel.tchannel import TChannel
+from tchannel.tcurl import tcurl
 
 
 @pytest.fixture
-def random_open_port():
-    """Find and return a random open TCP port."""
-    sock = socket.socket(socket.AF_INET)
-    try:
-        sock.bind(('', 0))
-        return sock.getsockname()[1]
-    finally:
-        sock.close()
-
-
-class Expectation(object):
-    """Represents an expectation for the ServerManager."""
-    def __init__(self, matcher):
-        assert matcher is not None
-
-        # expectation.matches(req) accepts a Message and returns True or
-        # False.
-        self.matches = matcher
-
-        # expectation.respond(context, connection) accepts the context and the
-        # connection and writes output to the connection.
-        self._respond = None
-
-    @classmethod
-    def messageType(cls, msg_typ):
-        """Build an expectation that expects a mesasge with the given type."""
-        return cls(lambda msg: msg.message_type == msg_typ)
-
-    @property
-    def respond(self):
-        # Do nothing if an action setter wasn't called.
-        if self._respond:
-            return self._respond
-        else:
-            return (lambda ctx, conn: None)
-
-    def and_return(self, resp):
-        """Write the given Message as a response."""
-        def respond(ctx, conn):
-            return conn.frame_and_write(resp)
-        self._respond = respond
-
-
-class ServerManager(object):
-    """Provides a dynamically configurable TChannel server."""
-    TIMEOUT = 0.15
-
-    def __init__(self, port, timeout=None):
-        manager = self
-        self.port = port
-        self.timeout = timeout or self.TIMEOUT
-
-        class Handler(SocketServer.BaseRequestHandler):
-            def setup(self):
-                self.request.settimeout(manager.timeout)
-                self.tchan_conn = tchannel.SocketConnection(self.request)
-
-            def handle(self):
-                (host, port) = self.request.getsockname()
-                self.tchan_conn.await_handshake(headers={
-                    'host_port': '%s:%s' % (host, port),
-                    'process_name': 'tchannel_server-%s' % port
-                })
-                self.tchan_conn.handle_calls(manager.handle_call)
-
-        self.server = SocketServer.TCPServer(("", port), Handler)
-        self.thread = None
-        self._expectations = []
-
-    def expect_ping(self):
-        """Expect a Ping request.
-
-        Returns an Expectation to allow setting the response behavior."""
-        exp = Expectation.messageType(
-            tmessage.PingRequestMessage.message_type
-        )
-
-        self._expectations.append(exp)
-        return exp
-
-    def handle_call(self, context, connection):
-        for exp in self._expectations:
-            if exp.matches(context.message):
-                exp.respond(context, connection)
-
-    @contextmanager
-    def client_connection(self):
-        """Get an initiated Connection to this TChannel server."""
-        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.settimeout(self.timeout)
-        try:
-            (host, port) = ('localhost', self.port)
-            sock.connect((host, port))
-
-            conn = tchannel.SocketConnection(sock)
-            conn.initiate_handshake(headers={
-                'host_port': '%s:%s' % (host, port),
-                'process_name': 'tchannel_client-%s' % port
-            })
-            conn.await_handshake_reply()
-            yield conn
-        finally:
-            sock.close()
-
-    def start(self):
-        assert self.thread is None, 'server already started'
-        self.thread = threading.Thread(target=self.server.serve_forever)
-        self.thread.start()
-
-    def stop(self):
-        self.server.shutdown()
-
-    def __enter__(self):
-        self.start()
-        return self
-
-    def __exit__(self, *args):
-        self.stop()
-
-
-@pytest.yield_fixture
-def server_manager(random_open_port):
-    with ServerManager(random_open_port) as manager:
-        yield manager
-
-
-def test_ping_pong(server_manager):
+def call_response():
+    resp = tmessage.CallResponseMessage()
+    resp.flags = 0
+    resp.code = 0
+    resp.span_id = 0
+    resp.parent_id = 0
+    resp.trace_id = 0
+    resp.traceflags = 0
+    resp.headers = {}
+    resp.checksum_type = 0
+    resp.checksum = 0
+    resp.arg_1 = 'hello'
+    resp.arg_2 = ''
+    resp.arg_3 = 'world'
+    return resp
+
+
+def test_tcp_ping_pong(server_manager):
     with server_manager.client_connection() as conn:
         resp = tmessage.PingResponseMessage()
         server_manager.expect_ping().and_return(resp)
 
         for i in range(1000):
             conn.ping()
             assert resp == next(conn).message
+
+
+@pytest.mark.gen_test
+def test_tchannel_call_request(server_manager, call_response):
+    endpoint = 'tchannelpeertest'
+    call_response.arg_1 = endpoint
+
+    server_manager.expect_call_request(endpoint).and_return(call_response)
+
+    tchannel = TChannel()
+
+    hostport = 'localhost:%d' % (server_manager.port)
+
+    response = yield tchannel.request(hostport).send(endpoint, '', '')
+
+    assert response.arg_1 == call_response.arg_1
+    assert response.arg_3 == call_response.arg_3
+
+
+@pytest.mark.gen_test
+def test_tcurl(server_manager, call_response):
+    endpoint = 'tcurltest'
+    call_response.arg_1 = endpoint
+
+    server_manager.expect_call_request(endpoint).and_return(call_response)
+
+    hostport = 'localhost:%d' % (server_manager.port)
+
+    response = yield tcurl(hostport, endpoint, '', '')
+
+    assert response.arg_1 == call_response.arg_1
+    assert response.arg_3 == call_response.arg_3

tests/integration/server_manager.py

@@ -0,0 +1,50 @@
+from __future__ import absolute_import
+
+import socket
+
+import pytest
+import tornado.gen
+import tornado.ioloop
+import tornado.testing
+
+from tchannel.tchannel import TChannel
+from tchannel.tornado.connection import TornadoConnection
+
+
+@pytest.mark.gen_test
+def test_add_peer_caching():
+    "Connections are long-lived and should not be recreated."""
+    tchannel = TChannel()
+    tchannel.peers = {'foo': 'bar'}
+    result = yield tchannel.add_peer('foo')
+    assert result == 'bar'
+
+
+def test_remove_peer():
+    tchannel = TChannel()
+    tchannel.peers = {'foo': 'bar'}
+    assert tchannel.remove_peer('foo') == 'bar'
+
+
+@pytest.mark.gen_test
+def test_get_peer_with_caching():
+    tchannel = TChannel()
+    tchannel.peers = {'foo': 'bar'}
+    result = yield tchannel.get_peer('foo')
+    assert result == 'bar'
+
+
+#@pytest.mark.gen_test
+#def test_tchannel_make_out_connection():
+    #server_sock, client_sock = socket.socketpair()
+
+    #server_stream = tornado.iostream.IOStream(server_sock)
+    #client_stream = tornado.iostream.IOStream(client_sock)
+
+    #server_conn = TornadoConnection(server_stream)
+    #client_conn = TornadoConnection(client_stream)
+
+    #hostname = server_sock.getsockname()
+
+    #tchannel = TChannel()
+    #conn = yield tchannel.make_out_connection(':', sock=server_sock)

tests/integration/test_tcp_client_server.py

@@ -0,0 +1,27 @@
+from __future__ import absolute_import
+
+import pytest
+import tornado.gen
+import tornado.ioloop
+import tornado.testing
+
+from tchannel.tornado.timeout import timeout
+
+
+class TimeoutTestCase(tornado.testing.AsyncTestCase):
+
+    @pytest.fixture(autouse=True)
+    def make_server_client(self, tornado_pair):
+        self.server, self.client = tornado_pair
+
+    def get_new_ioloop(self):
+        return tornado.ioloop.IOLoop.instance()
+
+    @tornado.testing.gen_test
+    def test_server_timeout(self):
+        with timeout(self.client, seconds=0.001):
+            future = self.client.initiate_handshake(headers={})
+            yield tornado.gen.sleep(0.001)
+            yield future
+
+        assert self.client.closed