[tor-commits] [stem/master] Adding a get_circuit() Controller method

[atagar at torproject.org]

commit cb9b4b605e8bdb639b22b2e58e5d2e80da9dcb56
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Dec 30 14:52:41 2012 -0800

    Adding a get_circuit() Controller method
    
    While reading through the tests I realized that we frequently wanted a method
    for 'gimme the circuit with id X', so adding a method for that.
---
 stem/control.py                  |   27 +++++++++++++++++++++++++++
 test/integ/control/controller.py |   15 +++++++--------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index 3c5f8bf..4d21849 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -36,6 +36,7 @@ providing its own for interacting at a higher level.
     |- close_stream - close a stream
     |- repurpose_circuit - change a circuit's purpose
     |- map_address - maps one address to another such that connections to the original are replaced with the other
+    |- get_circuit - provides an active circuit
     |- get_circuits - provides a list of active circuits
     |- attach_stream - attach a stream to a circuit
     |- get_version - convenience method to get tor version
@@ -1541,6 +1542,32 @@ class Controller(BaseController):
     
     return new_circuit
   
+  def get_circuit(self, circuit_id, default = UNDEFINED):
+    """
+    Provides a circuit presently available from tor.
+    
+    :param int circuit_id: circuit to be fetched
+    :param object default: response if the query fails
+    
+    :returns: :class:`stem.events.CircuitEvent` for the given circuit
+    
+    :raises:
+      * :class:`stem.ControllerError` if the call fails
+      * ValueError if the circuit doesn't exist
+      
+      An exception is only raised if we weren't provided a default response.
+    """
+    
+    try:
+      for circ in self.get_circuits():
+        if circ.id == circuit_id:
+          return circ
+      
+      raise ValueError("Tor presently does not have a circuit with the id of '%s'" % circuit_id)
+    except Exception, exc:
+      if default: return default
+      else: raise exc
+  
   def get_circuits(self):
     """
     Provides the list of circuits Tor is currently handling.
diff --git a/test/integ/control/controller.py b/test/integ/control/controller.py
index 74173c4..2dac687 100644
--- a/test/integ/control/controller.py
+++ b/test/integ/control/controller.py
@@ -520,10 +520,11 @@ class TestController(unittest.TestCase):
     
     with test.runner.get_runner().get_tor_controller() as controller:
       circuit_id = controller.extend_circuit('0')
+      
       # check if our circuit was created
-      self.assertTrue(filter(lambda circ: circ.id == circuit_id, controller.get_circuits()))
+      self.assertIsNotNone(controller.get_circuit(circuit_id, None))
       circuit_id = controller.new_circuit()
-      self.assertTrue(filter(lambda circ: circ.id == circuit_id, controller.get_circuits()))
+      self.assertIsNotNone(controller.get_circuit(circuit_id, None))
       
       self.assertRaises(stem.InvalidRequest, controller.extend_circuit, "foo")
       self.assertRaises(stem.InvalidRequest, controller.extend_circuit, '0', "thisroutershouldntexistbecausestemexists!@##$%#")
@@ -542,13 +543,11 @@ class TestController(unittest.TestCase):
     with runner.get_tor_controller() as controller:
       circ_id = controller.new_circuit()
       controller.repurpose_circuit(circ_id, "CONTROLLER")
-      circuits = controller.get_circuits()
-      circuit = filter(lambda circ: circ.id == circ_id, circuits)[0]
+      circuit = controller.get_circuit(circ_id)
       self.assertTrue(circuit.purpose == "CONTROLLER")
       
       controller.repurpose_circuit(circ_id, "GENERAL")
-      circuits = controller.get_circuits()
-      circuit = filter(lambda circ: circ.id == circ_id, circuits)[0]
+      circuit = controller.get_circuit(circ_id)
       self.assertTrue(circuit.purpose == "GENERAL")
       
       self.assertRaises(stem.InvalidRequest, controller.repurpose_circuit, 'f934h9f3h4', "fooo")
@@ -751,9 +750,9 @@ class TestController(unittest.TestCase):
       
       controller.add_event_listener(handle_streamcreated, stem.control.EventType.STREAM)
       ip = test.util.external_ip('127.0.0.1', socksport)
-      exit_circuit = [c for c in controller.get_circuits() if c.id == circuit_id]
+      exit_circuit = controller.get_circuit(circuit_id)
       self.assertTrue(exit_circuit)
-      exit_ip = controller.get_network_status(exit_circuit[0].path[2][0]).address
+      exit_ip = controller.get_network_status(exit_circuit.path[2][0]).address
       
       self.assertEquals(exit_ip, ip)