1 files changed, 35 insertions, 3 deletions

diff --git a/tests/integration/harness.py b/tests/integration/harness.py
index 868075d..7e75574 100644
--- a/tests/integration/harness.py
+++ b/tests/integration/harness.py
@@ -60,17 +60,20 @@ def free_port():
 class Tracker:
     """One tracker process, listening on loopback."""
 
-    def __init__(self, extra_config=""):
+    def __init__(self, extra_config="", files=None):
         self.port = free_port()
         self.dir = tempfile.mkdtemp(prefix="ot-integration-")
         self.log = None
         self.process = None
         try:
+            for name, content in (files or {}).items():
+                with open(os.path.join(self.dir, name), "wb") as handle:
+                    handle.write(content)
             self.config = os.path.join(self.dir, "opentracker.conf")
             with open(self.config, "w") as handle:
                 handle.write("listen.tcp 127.0.0.1:%d\n" % self.port)
                 handle.write("listen.udp 127.0.0.1:%d\n" % self.port)
-                handle.write(extra_config)
+                handle.write(extra_config.replace("{dir}", self.dir))
             self.log = open(os.path.join(self.dir, "tracker.log"), "w+b")
             # Leak detection is left at the sanitizer's own default, which is on
             # where LeakSanitizer exists and off where it does not. Asking for it
@@ -154,6 +157,12 @@ class Tracker:
     def stats(self, mode):
         return self.get("/stats?mode=" + mode)
 
+    def get_with_headers(self, path, headers=None, timeout=10):
+        """The body and the response headers, for the compressed answers."""
+        request = urllib.request.Request("http://127.0.0.1:%d%s" % (self.port, path), headers=headers or {})
+        with urllib.request.urlopen(request, timeout=timeout) as response:
+            return response.read(), dict(response.headers)
+
     def torrent_count(self):
         return int(self.stats("torr").split()[0])
 
@@ -180,14 +189,37 @@ class Tracker:
         assert action == 1, "udp announce answered action %d" % action
         return {"interval": interval, "leechers": leechers, "seeders": seeders, "peers": answer[20:]}
 
+    def udp_scrape(self, sock, connection_id, hashes):
+        """One (seeders, completed, leechers) triple per hash, in order."""
+        sock.sendto(struct.pack(">QII", connection_id, 2, 3) + b"".join(hashes), ("127.0.0.1", self.port))
+        answer = sock.recv(2048)
+        action, _ = struct.unpack(">II", answer[:8])
+        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 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."""
+        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        sock.settimeout(timeout)
+        try:
+            sock.sendto(payload, ("127.0.0.1", self.port))
+            try:
+                return sock.recv(2048)
+            except socket.timeout:
+                return None
+        finally:
+            sock.close()
+
 
 class TrackerTestCase(unittest.TestCase):
     """Gives every test its own tracker and fails it on any sanitizer output."""
 
     config = ""
+    files = None
 
     def setUp(self):
-        self.tracker = Tracker(self.config)
+        self.tracker = Tracker(self.config, self.files)
 
     def tearDown(self):
         complaints = self.tracker.stop()