Added testcases for keep-alive handling

* Added a dummy_server module to check connection count
* Added testcases for proxy and direct connections

diff --git a/tests/dummy_server.py b/tests/dummy_server.py
new file mode 100644 (file)
index 0000000..1096e25
--- /dev/null
+++ b/tests/dummy_server.py
@@ -0,0 +1,46 @@
+import asyncore\r
+import threading\r
+import socket\r
+\r
+class HttpServer(threading.Thread):\r
+    def __init__(self, port):\r
+        threading.Thread.__init__(self)\r
+        self.dispatcher = HttpServerDispatcher(port)\r
+\r
+    def run(self):\r
+        asyncore.loop()\r
+\r
+    @property\r
+    def connection_count(self):\r
+        return self.dispatcher.connection_count\r
+\r
+    def close(self):\r
+        asyncore.close_all()\r
+\r
+class HttpServerDispatcher(asyncore.dispatcher):\r
+    def __init__(self, port):\r
+        asyncore.dispatcher.__init__(self)\r
+        self.connected = False\r
+        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)\r
+        self.bind(('127.0.0.1', port))\r
+        self.listen(1)\r
+        self.connection_count = 0\r
+\r
+    def handle_accept(self):\r
+        self.connection_count += 1\r
+        self.handler = RequestHandler(self.accept()[0])\r
+\r
+    def handle_close(self):\r
+        self.close()\r
+\r
+\r
+class RequestHandler(asyncore.dispatcher_with_send):\r
+    def __init__(self, sock):\r
+        asyncore.dispatcher_with_send.__init__(self, sock)\r
+        self.response = ("HTTP/1.1 200 OK\r\n"\r
+                         "Connection: keep-alive\r\n"\r
+                         "Content-Length: 0\r\n\r\n")\r
+\r
+    def handle_read(self):\r
+        self.recv(1024)\r
+        self.send(self.response)\r

diff --git a/tests/test_keep_alive.py b/tests/test_keep_alive.py
new file mode 100644 (file)
index 0000000..447fc8e
--- /dev/null
+++ b/tests/test_keep_alive.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python\r
+# -*- coding: utf-8 -*-\r
+\r
+import os\r
+import sys\r
+import unittest\r
+\r
+# Path hack.\r
+sys.path.insert(0, os.path.abspath('..'))\r
+import requests\r
+import dummy_server\r
+\r
+class KeepAliveTests(unittest.TestCase):\r
+    server_and_proxy_port = 1234\r
+    request_count = 2\r
+    url = 'http://localhost:{0}'.format(server_and_proxy_port)\r
+    proxies={'http': url}\r
+\r
+    def setUp(self):\r
+        self.session = requests.session()\r
+        self.proxy_server = dummy_server.HttpServer(self.server_and_proxy_port)\r
+        self.proxy_server.start()\r
+\r
+    def tearDown(self):\r
+        self.proxy_server.close()\r
+\r
+    def test_keep_alive_with_direct_connection(self):\r
+        self.make_requests()\r
+        self.check_each_request_are_in_same_connection()\r
+\r
+    def test_no_keep_alive_with_direct_connection(self):\r
+        self.disable_keep_alive_in_session()\r
+        self.make_requests()\r
+        self.check_each_request_are_in_different_connection()\r
+\r
+    def test_keep_alive_with_proxy_connection(self):\r
+        self.make_proxy_requests()\r
+        self.check_each_request_are_in_same_connection()\r
+\r
+    def test_no_keep_alive_with_proxy_connection(self):\r
+        self.disable_keep_alive_in_session()\r
+        self.make_proxy_requests()\r
+        self.check_each_request_are_in_different_connection()\r
+\r
+    def make_proxy_requests(self):\r
+        self.make_requests(self.proxies)\r
+\r
+    def make_requests(self, proxies=None):\r
+        for _ in xrange(self.request_count):\r
+            self.session.get(self.url, proxies=proxies).text\r
+\r
+    def check_each_request_are_in_same_connection(self):\r
+        """Keep-alive requests open a single connection to the server."""\r
+        self.assertEqual(self.proxy_server.connection_count, 1)\r
+\r
+    def check_each_request_are_in_different_connection(self):\r
+        """Keep-alive requests open a single connection to the server."""\r
+        self.assertEqual(self.proxy_server.connection_count, self.request_count)\r
+\r
+    def disable_keep_alive_in_session(self):\r
+        self.session.config['keep_alive'] = False\r
+\r
+\r
+if __name__ == '__main__':\r
+    unittest.main()\r