[tor-commits] [stem/master] Adding get_effective_rate() to the Controller

[atagar at torproject.org]

commit 8ae878d48268ff38d232cc7c8e6ff40f8ec6324b
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Sep 20 13:12:51 2014 -0700

    Adding get_effective_rate() to the Controller
    
    Method to provide the bytes per second we're configured to relay at. Tor has
    this be the minimum of a few parameters so it's not as stright forward as
    'gimme the GETCONF BandwidthRate'.
---
 docs/change_log.rst             |    1 +
 stem/control.py                 |   42 +++++++++++++++++++++++++++++++++++++++
 test/unit/control/controller.py |   23 +++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index c5867e4..7297cb9 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -43,6 +43,7 @@ The following are only available within Stem's `git repository
  * **Controller**
 
   * Added :func:`~stem.control.Controller.get_accounting_stats` to the :class:`~stem.control.Controller`
+  * Added :func:`~stem.control.Controller.get_effective_rate` to the :class:`~stem.control.Controller`
   * Added :func:`~stem.control.BaseController.connection_time` to the :class:`~stem.control.BaseController`
   * Changed :func:`~stem.control.Controller.get_microdescriptor`, :func:`~stem.control.Controller.get_server_descriptor`, and :func:`~stem.control.Controller.get_network_status` to get our own descriptor if no fingerprint or nickname is provided.
   * Added :class:`~stem.exit_policy.ExitPolicy` methods for more easily handling 'private' policies (the `default prefix <https://www.torproject.org/docs/tor-manual.html.en#ExitPolicyRejectPrivate>`_) and the defaultly appended suffix. This includes :func:`~stem.exit_policy.ExitPolicy.has_private`, :func:`~stem.exit_policy.ExitPolicy.strip_private`, :func:`~stem.exit_policy.ExitPolicy.has_default`, and :func:`~stem.exit_policy.ExitPolicy.strip_default` :class:`~stem.exit_policy.ExitPolicy` methods in addition to :func:`~stem.exit_policy.ExitPolicyRule.is_private` and :func:`~stem.exit_policy.ExitPolicyRule.is_default` for the :class:`~stem.exit_policy.ExitPolicyRule`. (:trac:`10107`)
diff --git a/stem/control.py b/stem/control.py
index 04098c0..4688bd7 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -122,6 +122,7 @@ If you're fine with allowing your script to raise exceptions then this can be mo
     |- signal - sends a signal to the tor client
     |- is_newnym_available - true if tor would presently accept a NEWNYM signal
     |- get_newnym_wait - seconds until tor would accept a NEWNYM signal
+    |- get_effective_rate - provides our effective relaying rate limit
     |- is_geoip_unavailable - true if we've discovered our geoip db to be unavailable
     |- map_address - maps one address to another such that connections to the original are replaced with the other
     +- drop_guards - drops our set of guard relays and picks a new set
@@ -2666,6 +2667,47 @@ class Controller(BaseController):
 
     return max(0.0, self._last_newnym + 10 - time.time())
 
+  def get_effective_rate(self, default = UNDEFINED, burst = False):
+    """
+    Provides the maximum rate this relay is configured to relay in bytes per
+    second. This is based on multiple torrc parameters if they're set...
+
+    * Effective Rate = min(BandwidthRate, RelayBandwidthRate, MaxAdvertisedBandwidth)
+    * Effective Burst = min(BandwidthBurst, RelayBandwidthBurst)
+
+    :param object default: response if the query fails
+    :param bool burst: provides the burst bandwidth, otherwise this provides
+      the standard rate
+
+    :returns: **int** with the effective bandwidth rate in bytes per second
+
+    :raises: :class:`stem.ControllerError` if the call fails and no default was
+      provided
+    """
+
+    if not burst:
+      attributes = ('BandwidthRate', 'RelayBandwidthRate', 'MaxAdvertisedBandwidth')
+    else:
+      attributes = ('BandwidthBurst', 'RelayBandwidthBurst')
+
+    value = None
+
+    for attr in attributes:
+      try:
+        attr_value = int(self.get_conf(attr))
+
+        if attr_value == 0 and attr.startswith('Relay'):
+          continue  # RelayBandwidthRate and RelayBandwidthBurst default to zero
+
+        value = min(value, attr_value) if value else attr_value
+      except stem.ControllerError as exc:
+        if default == UNDEFINED:
+          raise exc
+        else:
+          return default
+
+    return value
+
   def is_geoip_unavailable(self):
     """
     Provides **True** if we've concluded hat our geoip database is unavailable,
diff --git a/test/unit/control/controller.py b/test/unit/control/controller.py
index 7c7911e..948068f 100644
--- a/test/unit/control/controller.py
+++ b/test/unit/control/controller.py
@@ -588,6 +588,29 @@ class TestControl(unittest.TestCase):
     for test_input in malformed_inputs:
       self.assertRaises(ProtocolError, _parse_circ_path, test_input)
 
+  @patch('stem.control.Controller.get_conf')
+  def test_get_effective_rate(self, get_conf_mock):
+    """
+    Exercise the get_effective_rate() method.
+    """
+
+    # check default if nothing was set
+
+    get_conf_mock.side_effect = lambda param, **kwargs: {
+      'BandwidthRate': '1073741824',
+      'BandwidthBurst': '1073741824',
+      'RelayBandwidthRate': '0',
+      'RelayBandwidthBurst': '0',
+      'MaxAdvertisedBandwidth': '1073741824',
+    }[param]
+
+    self.assertEqual(1073741824, self.controller.get_effective_rate())
+    self.assertEqual(1073741824, self.controller.get_effective_rate(burst = True))
+
+    get_conf_mock.side_effect = ControllerError('nope, too bad')
+    self.assertRaises(ControllerError, self.controller.get_effective_rate)
+    self.assertEqual('my_default', self.controller.get_effective_rate('my_default'))
+
   @patch('stem.control.Controller.get_version')
   def test_drop_guards(self, get_version_mock):
     """