Add new test for connection keep alive

1 files changed, 78 insertions, 0 deletions

diff --git a/tests/integration/harness.py b/tests/integration/harness.py
index 7e75574..03a0561 100644
--- a/tests/integration/harness.py
+++ b/tests/integration/harness.py
@@ -197,6 +197,10 @@ class Tracker:
         assert action == 2, "udp scrape answered action %d" % action
         return [struct.unpack(">III", answer[8 + 12 * i : 20 + 12 * i]) for i in range(len(hashes))]
 
+    def connect(self):
+        """A connection the caller drives and closes."""
+        return Connection(self.port)
+
     def udp_raw(self, payload, timeout=1.0):
         """Sends a packet exactly as given. Returns the answer, or None if the
         tracker chose to say nothing, which is how it treats a broken one."""
@@ -212,6 +216,73 @@ class Tracker:
             sock.close()
 
 
+class Connection:
+    """One HTTP connection the caller keeps open, which urllib will not do.
+
+    Reads exactly one reply at a time, so a test can tell "the tracker
+    answered and kept the socket" from "the tracker answered and hung up"."""
+
+    def __init__(self, port, timeout=5):
+        self.sock = socket.create_connection(("127.0.0.1", port), timeout=timeout)
+        self.sock.settimeout(timeout)
+        self.buffer = b""
+
+    def send(self, path, connection="Keep-Alive", version="1.1"):
+        request = "GET %s HTTP/%s\r\nHost: tracker\r\n" % (path, version)
+        if connection is not None:
+            request += "Connection: %s\r\n" % connection
+        self.sock.sendall((request + "\r\n").encode())
+
+    def _fill(self):
+        chunk = self.sock.recv(65536)
+        if not chunk:
+            raise EOFError("the tracker closed the connection")
+        self.buffer += chunk
+
+    def reply(self):
+        """The one complete reply at the front of the stream, headers and body."""
+        while b"\r\n\r\n" not in self.buffer:
+            self._fill()
+        head, rest = self.buffer.split(b"\r\n\r\n", 1)
+        headers = {}
+        for line in head.split(b"\r\n")[1:]:
+            name, _, value = line.partition(b":")
+            headers[name.strip().lower()] = value.strip()
+
+        if headers.get(b"transfer-encoding") == b"chunked":
+            while not rest.endswith(b"0\r\n\r\n"):
+                self.buffer = head + b"\r\n\r\n" + rest
+                self._fill()
+                rest = self.buffer.split(b"\r\n\r\n", 1)[1]
+            body, self.buffer = rest, b""
+        else:
+            length = int(headers.get(b"content-length", b"0"))
+            while len(rest) < length:
+                self.buffer = head + b"\r\n\r\n" + rest
+                self._fill()
+                rest = self.buffer.split(b"\r\n\r\n", 1)[1]
+            body, self.buffer = rest[:length], rest[length:]
+        return head.split(b"\r\n")[0], headers, body
+
+    def is_open(self):
+        """True if the tracker is still holding the connection."""
+        try:
+            self.sock.settimeout(0.4)
+            return self.sock.recv(1) != b""
+        except socket.timeout:
+            return True
+        except OSError:
+            return False
+        finally:
+            try:
+                self.sock.settimeout(5)
+            except OSError:
+                pass
+
+    def close(self):
+        self.sock.close()
+
+
 class TrackerTestCase(unittest.TestCase):
     """Gives every test its own tracker and fails it on any sanitizer output."""
 
@@ -222,6 +293,13 @@ class TrackerTestCase(unittest.TestCase):
         self.tracker = Tracker(self.config, self.files)
 
     def tearDown(self):
+        # A connection the test closed is only reaped when the event loop gets
+        # round to it, and a cookie still held at exit reads as a leak. One
+        # more request puts us behind that work in the loop's queue.
+        try:
+            self.tracker.torrent_count()
+        except Exception:
+            pass
         complaints = self.tracker.stop()
         self.assertEqual(complaints, [], "the tracker reported:\n" + "\n".join(complaints))