[tor-commits] [stem/master] Add a Controller.get_streams method and its tests

[atagar at torproject.org]

commit 7e5377906b6b0e0bcc20fb442fb11554df929868
Author: Sean Robinson <seankrobinson at gmail.com>
Date:   Wed Jan 2 05:52:29 2013 -0700

    Add a Controller.get_streams method and its tests
    
    Re-use Ravi's idea for parsing get_info("circuit-status") responses with
    get_info("stream-status").  The stream status GETINFO reply contains just
    the mandatory parts of a STREAM event, so the event optional keyword
    arguments are not (yet) used.
    
    Signed-off-by: Sean Robinson <seankrobinson at gmail.com>
---
 stem/control.py                  |   20 ++++++++++++++++++++
 test/integ/control/controller.py |   25 +++++++++++++++++++++++++
 test/unit/control/controller.py  |   31 ++++++++++++++++++++++++++++++-
 3 files changed, 75 insertions(+), 1 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index d8ee940..8bb3a10 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -51,6 +51,7 @@ providing its own for interacting at a higher level.
     |- close_circuit - close a circuit
     |
     |- attach_stream - attach a stream to a circuit
+    |- get_streams - provides a list of active streams
     |- close_stream - close a stream
     |
     |- signal - sends a signal to the tor client
@@ -1716,6 +1717,25 @@ class Controller(BaseController):
       else:
         raise stem.ProtocolError("ATTACHSTREAM returned unexpected response code: %s" % response.code)
   
+  def get_streams(self):
+    """
+    Provides the list of streams tor is currently handling.
+    
+    :returns: list of :class:`stem.events.StreamEvent` objects
+    
+    :raises: :class:`stem.ControllerError` if the call fails
+    """
+    
+    streams = []
+    response = self.get_info("stream-status")
+    
+    for stream in response.splitlines():
+      message = stem.socket.recv_message(StringIO.StringIO("650 STREAM " + stream + "\r\n"))
+      stem.response.convert("EVENT", message, arrived_at = 0)
+      streams.append(message)
+    
+    return streams
+  
   def close_stream(self, stream_id, reason = stem.RelayEndReason.MISC, flag = ''):
     """
     Closes the specified stream.
diff --git a/test/integ/control/controller.py b/test/integ/control/controller.py
index 534bba2..25fc431 100644
--- a/test/integ/control/controller.py
+++ b/test/integ/control/controller.py
@@ -20,6 +20,7 @@ import stem.descriptor.router_status_entry
 import stem.response.protocolinfo
 import stem.socket
 import stem.version
+import test.network
 import test.runner
 import test.util
 
@@ -580,6 +581,30 @@ class TestController(unittest.TestCase):
       self.assertRaises(stem.InvalidArguments, controller.close_circuit, str(int(circ_id) + 1024))
       self.assertRaises(stem.InvalidRequest, controller.close_circuit, "")
   
+  def test_get_streams(self):
+    """
+    Tests Controller.get_streams().
+    """
+    
+    if test.runner.require_control(self): return
+    elif test.runner.require_online(self): return
+    
+    host = "38.229.72.14"   # www.torproject.org
+    port = 443
+    target = host + ":" + str(port)
+    
+    runner = test.runner.get_runner()
+    with runner.get_tor_controller() as controller:
+      # we only need one proxy port, so take the first
+      socks_listener = controller.get_socks_listeners()[0]
+      with test.network.Socks(socks_listener) as s:
+        s.settimeout(30)
+        s.connect((host, port))
+        streams = controller.get_streams()
+    # Because we do not get a stream id when opening a stream,
+    #  try to match the target for which we asked a stream.
+    self.assertTrue(target in [stream.target for stream in streams])
+  
   def test_mapaddress(self):
     if test.runner.require_control(self): return
     elif test.runner.require_online(self): return
diff --git a/test/unit/control/controller.py b/test/unit/control/controller.py
index 695c2de..a6fc372 100644
--- a/test/unit/control/controller.py
+++ b/test/unit/control/controller.py
@@ -10,6 +10,7 @@ import stem.version
 
 from stem import InvalidArguments, InvalidRequest, ProtocolError
 from stem.control import _parse_circ_path, Controller, EventType
+from stem.response import events
 from test import mocking
 
 class TestControl(unittest.TestCase):
@@ -167,4 +168,32 @@ class TestControl(unittest.TestCase):
     for response in invalid_responses:
       mocking.mock_method(Controller, "get_info", mocking.return_value(response))
       self.assertRaises(stem.ProtocolError, self.controller.get_socks_listeners)
-
+  
+  def test_get_streams(self):
+    """
+    Exercises the get_streams() method.
+    """
+    
+    # get a list of fake, but good looking, streams
+    valid_streams = (
+      ("1", "NEW", "4", "10.10.10.1:80"),
+      ("2", "SUCCEEDED", "4", "10.10.10.1:80"),
+      ("3", "SUCCEEDED", "4", "10.10.10.1:80")
+    )
+    response = []
+    for stream_parts in valid_streams:
+      stream = mocking.get_object(events.StreamEvent)
+      (stream.id, stream.status, stream.circ_id, stream.target) = stream_parts
+      line = "%s\r\n" % " ".join(stream_parts)
+      response.append(line)
+    
+    mocking.mock_method(Controller, "get_info", mocking.return_value(
+      "".join(response)
+    ))
+    streams = self.controller.get_streams()
+    self.assertEqual(len(valid_streams), len(streams))
+    for index in range(len(streams)):
+      self.assertEqual(valid_streams[index][0], streams[index].id)
+      self.assertEqual(valid_streams[index][1], streams[index].status)
+      self.assertEqual(valid_streams[index][2], streams[index].circ_id)
+      self.assertEqual(valid_streams[index][3], streams[index].target)