[tor-commits] [arm/master] Testing the ConnectionTracker class

[atagar at torproject.org]

commit 6f2ce6c6d1004fd69d8ba1038f08c1376cdf2625
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Jan 5 17:55:17 2014 -0800

    Testing the ConnectionTracker class
    
    Unit tests for our connection resolver daemon. There's a lot more tests we
    could add for the tracker module, but this gets us pretty good coverage for
    now.
---
 arm/util/tracker.py                     |    2 +-
 test/util/tracker/connection_tracker.py |   76 +++++++++++++++++++++++++++++++
 test/util/tracker/resource_tracker.py   |    5 +-
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/arm/util/tracker.py b/arm/util/tracker.py
index b0f99b2..6c5b1f4 100644
--- a/arm/util/tracker.py
+++ b/arm/util/tracker.py
@@ -384,7 +384,7 @@ class ConnectionTracker(Daemon):
         self._failure_count += 1
 
         if self._failure_count >= 3:
-          self._resolvers.pop()
+          self._resolvers.pop(0)
           self._failure_count = 0
 
           if self._resolvers:
diff --git a/test/util/tracker/connection_tracker.py b/test/util/tracker/connection_tracker.py
new file mode 100644
index 0000000..bdb1d69
--- /dev/null
+++ b/test/util/tracker/connection_tracker.py
@@ -0,0 +1,76 @@
+import time
+import unittest
+
+from arm.util.tracker import ConnectionTracker
+
+from stem.util import connection
+
+from mock import Mock, patch
+
+CONNECTION_1 = connection.Connection('127.0.0.1', 3531, '75.119.206.243', 22, 'tcp')
+CONNECTION_2 = connection.Connection('127.0.0.1', 1766, '86.59.30.40', 443, 'tcp')
+CONNECTION_3 = connection.Connection('127.0.0.1', 1059, '74.125.28.106', 80, 'tcp')
+
+
+class TestConnectionTracker(unittest.TestCase):
+  @patch('arm.util.tracker.tor_controller')
+  @patch('arm.util.tracker.connection.get_connections')
+  @patch('arm.util.tracker.system', Mock(return_value = Mock()))
+  @patch('arm.util.tracker.connection.get_system_resolvers', Mock(return_value = [connection.Resolver.NETSTAT]))
+  def test_fetching_connections(self, get_connections_mock, tor_controller_mock):
+    tor_controller_mock().get_pid.return_value = 12345
+    get_connections_mock.return_value = [CONNECTION_1, CONNECTION_2, CONNECTION_3]
+
+    with ConnectionTracker(0.04) as daemon:
+      time.sleep(0.01)
+
+      connections = daemon.get_connections()
+
+      self.assertEqual(1, daemon.run_counter())
+      self.assertEqual([CONNECTION_1, CONNECTION_2, CONNECTION_3], connections)
+
+      get_connections_mock.return_value = []  # no connection results
+      time.sleep(0.05)
+      connections = daemon.get_connections()
+
+      self.assertEqual(2, daemon.run_counter())
+      self.assertEqual([], connections)
+
+  @patch('arm.util.tracker.tor_controller')
+  @patch('arm.util.tracker.connection.get_connections')
+  @patch('arm.util.tracker.system', Mock(return_value = Mock()))
+  @patch('arm.util.tracker.connection.get_system_resolvers', Mock(return_value = [connection.Resolver.NETSTAT, connection.Resolver.LSOF]))
+  def test_resolver_failover(self, get_connections_mock, tor_controller_mock):
+    tor_controller_mock().get_pid.return_value = 12345
+    get_connections_mock.side_effect = IOError()
+
+    with ConnectionTracker(0.01) as daemon:
+      time.sleep(0.03)
+
+      self.assertEqual([connection.Resolver.NETSTAT, connection.Resolver.LSOF], daemon._resolvers)
+      self.assertEqual([], daemon.get_connections())
+
+      time.sleep(0.05)
+
+      self.assertEqual([connection.Resolver.LSOF], daemon._resolvers)
+      self.assertEqual([], daemon.get_connections())
+
+      time.sleep(0.05)
+
+      self.assertEqual([], daemon._resolvers)
+      self.assertEqual([], daemon.get_connections())
+
+      # Now make connection resolution work. We still shouldn't provide any
+      # results since we stopped looking.
+
+      get_connections_mock.return_value = [CONNECTION_1, CONNECTION_2]
+      get_connections_mock.side_effect = None
+      time.sleep(0.05)
+      self.assertEqual([], daemon.get_connections())
+
+      # Finally, select a custom resolver. This should cause us to query again
+      # reguardless of our prior failures.
+
+      daemon.set_custom_resolver(connection.Resolver.NETSTAT)
+      time.sleep(0.05)
+      self.assertEqual([CONNECTION_1, CONNECTION_2], daemon.get_connections())
diff --git a/test/util/tracker/resource_tracker.py b/test/util/tracker/resource_tracker.py
index 70cdf27..41ef80e 100644
--- a/test/util/tracker/resource_tracker.py
+++ b/test/util/tracker/resource_tracker.py
@@ -107,12 +107,13 @@ class TestResourceTracker(unittest.TestCase):
       self.assertEqual(0.0, resources.memory_percent)
       self.assertEqual(0.0, resources.timestamp)
 
-      time.sleep(0.05)
+      while daemon.run_counter() < 1:
+        time.sleep(0.01)
+
       self.assertEqual(False, daemon._use_proc)
 
       resources = daemon.get_resource_usage()
 
-      self.assertEqual(1, daemon.run_counter())
       self.assertEqual(0.0, resources.cpu_sample)
       self.assertEqual(43.875, resources.cpu_average)
       self.assertEqual(105.3, resources.cpu_total)