Refactor tests to avoid duplicate path selection

Author: DonnchaC

bwscanner/circuit.py

@@ -7,28 +7,39 @@
 from bwscanner.logger import log
 
 
+def is_valid_exit(relay):
+    """
+    Check that has the correct flags and exit policy for exiting
+
+    TODO: Check the exit policy
+    """
+    is_exit = ('exit' in relay.flags and 'badexit' not in relay.flags)
+    return is_exit and 'authority' not in relay.flags
+
+
+def random_path_to_exit(exit_relay, relays):
+    """
+    Choose a random path to the specified exit relay
+    """
+    # Sample an extra relay in case one of the choices is the exit.
+    candidate_relays = random.sample(relays, 3)
+    if exit_relay in candidate_relays:
+        candidate_relays.remove(exit_relay)
+    return candidate_relays[0:2] + [exit_relay]
+
+
 class CircuitGenerator(object):
     def __init__(self, state):
         self.state = state
         self.relays = list(set(state.routers.values()))
-        self.exits = [relay for relay in self.relays if self.is_valid_exit(relay)]
+        self.exits = [relay for relay in self.relays if is_valid_exit(relay)]
 
     def __iter__(self):
         return self
 
     def next(self):
         raise NotImplementedError
 
-    @staticmethod
-    def is_valid_exit(relay):
-        """
-        Check that has the correct flags and exit policy for exiting
-
-        TODO: Check the exit policy
-        """
-        is_exit = ('exit' in relay.flags and 'badexit' not in relay.flags)
-        return is_exit and 'authority' not in relay.flags
-
 
 class ExitScan(CircuitGenerator):
     """
@@ -48,11 +59,7 @@ def circuit_generator():
             to use the exit relay more than once.
             """
             for exit_relay in self.exits:
-                # Sample an extra relay in case one of the choices is the exit.
-                candidate_relays = random.sample(self.relays, 3)
-                if exit_relay in candidate_relays:
-                    candidate_relays.remove(exit_relay)
-                yield candidate_relays[0:2] + [exit_relay]
+                yield random_path_to_exit(exit_relay, self.relays)
 
         self._circgen = circuit_generator()
 

test/template.py

@@ -6,14 +6,11 @@
 from twisted.trial import unittest
 from txtorcon.torstate import build_tor_connection
 
+from bwscanner import circuit
 from bwscanner.attacher import SOCKSClientStreamAttacher
 
 
 class TorTestCase(unittest.TestCase):
-    """
-    XXX: Use code from circuit.py for the path selection rather than
-         repeating path selection here.
-    """
 
     @defer.inlineCallbacks
     def setUp(self):
@@ -30,14 +27,11 @@ def routers(self):
 
     @property
     def exits(self):
-        return [r for r in self.routers if 'exit' in r.flags]
+        return [r for r in self.routers if circuit.is_valid_exit(r)]
 
     def random_path(self):
         exit_relay = random.choice(self.exits)
-        selected_relays = random.sample(self.routers, 3)
-        if exit_relay in selected_relays:
-            selected_relays.remove(exit_relay)
-        return selected_relays[0:2] + [exit_relay]
+        return circuit.random_path_to_exit(exit_relay, self.routers)
 
     @defer.inlineCallbacks
     def tearDown(self):