1 files changed, 20 insertions, 7 deletions
diff --git a/lib/chutney/Testing.py b/lib/chutney/Testing.py
index 8b1b356..341f519 100644
--- a/lib/chutney/Testing.py
+++ b/lib/chutney/Testing.py
@@ -240,14 +240,23 @@ def create_hidden_service(nodes):
         isLeaf = True
         numberRequests = 0
 
-        def __init__(self, siteNum):
+        def __init__(self, siteNum, port):
             self.siteNum = siteNum
+            self.port = port
 
         def render_GET(self, request):
             self.numberRequests += 1
             request.setHeader("content-type", "text/plain")
             return str(self.siteNum)
 
+        def listen(self):
+            site = server.Site(self)
+            self.s = reactor.listenTCP(self.port, site)
+
+        def stopListening(self):
+            self.s.stopListening()
+
+
     base_port = 10080
     for i in range(nodes):
 
@@ -261,10 +270,7 @@ def create_hidden_service(nodes):
         )
         hs_nodes.append(node)
 
-        site = server.Site(Site(i))
-
-        s = reactor.listenTCP(port, site)
-        hs_servers.append(s)
+        hs_servers.append(Site(i, port))
 
         port += 1
 
@@ -285,13 +291,20 @@ def connection_test(client_nodes, expected_nodes):
 
     test_pass = True
 
+    nodes_detected = 0
+
     for node, responses in sorted(responses.items()):
         logging.info(node + ": " + str(responses))
 
-        if responses == 0:
-            test_pass = False
+        if responses != 0:
+            nodes_detected += 1
+
+    if nodes_detected != expected_nodes:
+        test_pass = False
 
     if test_pass:
         logging.info("Test PASS")
     else:
         logging.info("Test FAIL")
+
+    return test_pass