import re
import socket

from tornado.concurrent import Future
from tornado.netutil import Resolver
from tornado.tcpserver import TCPServer
from tornado.testing import AsyncTestCase, AsyncHTTPTestCase
from tornado.testing import gen_test, bind_unused_port
from tornado.web import RequestHandler, Application

from trickle import Trickle


class TestTCPServer(TCPServer):
    def __init__(self, *args, **kwargs):
        super(TestTCPServer, self).__init__(*args, **kwargs)
        self.test_stream = Future()

    def handle_stream(self, stream, address):
        self.test_stream.set_result(stream)


class TrickleTCPTest(AsyncTestCase):
    # Tests with a TCPServer.

    def setUp(self):
        super(TrickleTCPTest, self).setUp()
        sock, port = bind_unused_port()
        self.port = port
        self.server = TestTCPServer(self.io_loop)
        self.server.add_socket(sock)
        self.resolver = Resolver()

    @gen_test
    def test_basic(self):
        client_trickle = Trickle(
            socket.socket(socket.AF_INET),
            io_loop=self.io_loop)

        addr_info = yield self.resolver.resolve(
            'localhost', self.port, socket.AF_INET)

        sock_addr = addr_info[0][1]
        yield client_trickle.connect(sock_addr)

        # Wait for server to handle connection.
        server_stream = yield self.server.test_stream
        server_trickle = Trickle(server_stream)
        data = b'a' * 10
        yield server_trickle.write(data)
        self.assertEqual(data, (yield client_trickle.read_bytes(10)))

    @gen_test
    def test_read_timeout(self):
        client_trickle = Trickle(
            socket.socket(socket.AF_INET),
            io_loop=self.io_loop)

        addr_info = yield self.resolver.resolve(
            'localhost', self.port, socket.AF_INET)

        sock_addr = addr_info[0][1]
        yield client_trickle.connect(sock_addr)

        try:
            yield client_trickle.read_bytes(10, timeout=0.01)
        except socket.timeout:
            pass
        else:
            self.fail('socket.timeout not raised')


class TrickleTestHandler(RequestHandler):
    def get(self):
        self.finish('hello')


class TrickleHTTPTest(AsyncHTTPTestCase):
    # Tests with an HTTPServer.

    def setUp(self):
        super(TrickleHTTPTest, self).setUp()
        self.resolver = Resolver()

    def get_app(self):
        return Application([('/', TrickleTestHandler)])

    @gen_test
    def test_http(self):
        addr_info = yield self.resolver.resolve(
            'localhost',
            self.get_http_port(),
            socket.AF_INET)

        sock_addr = addr_info[0][1]
        trick = Trickle(
            socket.socket(socket.AF_INET),
            io_loop=self.io_loop)

        yield trick.connect(sock_addr)
        yield trick.write(b'GET / HTTP/1.1\r\n\r\n')

        headers = yield trick.read_until(b'\r\n\r\n')
        match = re.search(br'Content-Length: (\d+)\r\n', headers)
        content_length = int(match.group(1))
        body = yield trick.read_bytes(content_length)
        self.assertEqual('hello', body)

@@ -18,16 +18,40 @@
 socket%0A%0A
+from tornado import gen%0A
 from tor
@@ -80,16 +80,16 @@
  Future%0A
-
 from tor
@@ -952,32 +952,34 @@
 @gen
-_test%0A    def test_basic
+.coroutine%0A    def connect
 (sel
@@ -1430,16 +1430,184 @@
 _stream)
+%0A%0A        raise gen.Return((client_trickle, server_trickle))%0A%0A    @gen_test%0A    def test_read_bytes(self):%0A        client_trickle, server_trickle = yield self.connect()
 %0A       
@@ -1809,282 +1809,46 @@
 ckle
- = Trickle(%0A            socket.socket(socket.AF_INET),%0A            io_loop=self.io_loop)%0A%0A        addr_info = yield self.resolver.resolve(%0A            'localhost', self.port, socket.AF_INET)%0A%0A        sock_addr = addr_info%5B0%5D%5B1%5D%0A        yield client_trickle.connect(sock_addr
+, server_trickle = yield self.connect(
 )%0A%0A 

Test refactoring.