[tor-commits] [stem/master] Test benchmark_stem example

[atagar at torproject.org]

commit 4497148d0be8b267e032d8867c082a3588529581
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Sep 21 17:26:30 2020 -0700

    Test benchmark_stem example
    
    Caught our first bug in our example code... of sorts. This adds a
    'descriptor_type' argument to benchmark_stem.py because CollecTor
    microdescriptor archives contain a mixture of consensuses and microdescriptors.
---
 docs/_static/example/benchmark_stem.py | 11 +++---
 test/unit/examples.py                  | 62 +++++++++++++++++++++++++++++++---
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/docs/_static/example/benchmark_stem.py b/docs/_static/example/benchmark_stem.py
index 8b71d730..7094c3dd 100644
--- a/docs/_static/example/benchmark_stem.py
+++ b/docs/_static/example/benchmark_stem.py
@@ -60,7 +60,7 @@ def measure_fraction_relays_exit_80_microdescriptors(path):
   start_time = time.time()
   exits, count = 0, 0
 
-  for desc in stem.descriptor.parse_file(path):
+  for desc in stem.descriptor.parse_file(path, descriptor_type = 'microdescriptor 1.0'):
     if desc.exit_policy.can_exit_to(port = 80):
       exits += 1
 
@@ -74,8 +74,9 @@ def measure_fraction_relays_exit_80_microdescriptors(path):
   print('  Time per microdescriptor: %0.5f seconds' % (runtime / count))
   print('')
 
-measure_average_advertised_bandwidth('/home/atagar/Desktop/server-descriptors-2015-11.tar')
-measure_countries_v3_requests('/home/atagar/Desktop/extra-infos-2015-11.tar')
-measure_average_relays_exit('/home/atagar/Desktop/consensuses-2015-11.tar')
-measure_fraction_relays_exit_80_microdescriptors('/home/atagar/Desktop/microdescs-2015-11.tar')
+if __name__ == '__main__':
+  measure_average_advertised_bandwidth('/home/atagar/Desktop/server-descriptors-2015-11.tar')
+  measure_countries_v3_requests('/home/atagar/Desktop/extra-infos-2015-11.tar')
+  measure_average_relays_exit('/home/atagar/Desktop/consensuses-2015-11.tar')
+  measure_fraction_relays_exit_80_microdescriptors('/home/atagar/Desktop/microdescs-2015-11.tar')
 
diff --git a/test/unit/examples.py b/test/unit/examples.py
index 6a00d9b7..124ca3c7 100644
--- a/test/unit/examples.py
+++ b/test/unit/examples.py
@@ -15,6 +15,7 @@ from stem.descriptor.bandwidth_file import BandwidthFile
 from unittest.mock import patch
 
 EXAMPLE_DIR = os.path.join(test.STEM_BASE, 'docs', '_static', 'example')
+DESC_DIR = os.path.join(test.STEM_BASE, 'test', 'unit', 'descriptor', 'data')
 
 EXPECTED_BANDWIDTH_STATS = """\
 Relay FDCF49562E65B1CC219410009BD48A9EED387C77
@@ -33,7 +34,7 @@ Relay BD4172533C3F7271ABCCD9F057E06FD91547C42B
 
 """
 
-EXPECTED_BENCHMARK_SERVER_DESC_PREFIX = """\
+EXPECTED_SERVER_DESC_BENCHMARK_PREFIX = """\
 Finished measure_average_advertised_bandwidth('%s')
   Total time: 0 seconds
   Processed server descriptors: 5
@@ -41,6 +42,30 @@ Finished measure_average_advertised_bandwidth('%s')
   Time per server descriptor:
 """.rstrip()
 
+EXPECTED_EXTRAINFO_BENCHMARK_PREFIX = """\
+Finished measure_countries_v3_requests('%s')
+  Total time: 0 seconds
+  Processed extrainfo descriptors: 7
+  Number of countries: 6
+  Time per extrainfo descriptor:
+""".rstrip()
+
+EXPECTED_CONSENSUS_BENCHMARK_PREFIX = """\
+Finished measure_average_relays_exit('%s')
+  Total time: 0 seconds
+  Processed 2 consensuses with 243 router status entries
+  Total exits: 28 (0.12%%)
+  Time per consensus:
+""".rstrip()
+
+EXPECTED_MICRODESC_BENCHMARK_PREFIX = """\
+Finished measure_fraction_relays_exit_80_microdescriptors('%s')
+  Total time: 0 seconds
+  Processed microdescriptors: 3
+  Total exits to port 80: 1 (0.33%%)
+  Time per microdescriptor:
+""".rstrip()
+
 
 def import_example(module_name):
   """
@@ -89,15 +114,44 @@ class TestExamples(unittest.TestCase):
 
   @patch('sys.stdout', new_callable = io.StringIO)
   def test_benchmark_server_descriptor_stem(self, stdout_mock):
-    path = os.path.join(test.STEM_BASE, 'test', 'unit', 'descriptor', 'data', 'collector', 'server-descriptors-2005-12-cropped.tar')
+    path = os.path.join(DESC_DIR, 'collector', 'server-descriptors-2005-12-cropped.tar')
+    expected_prefix = EXPECTED_SERVER_DESC_BENCHMARK_PREFIX % path
 
     module = import_example('benchmark_server_descriptor_stem')
     module.measure_average_advertised_bandwidth(path)
 
-    self.assertTrue(stdout_mock.getvalue().startswith(EXPECTED_BENCHMARK_SERVER_DESC_PREFIX % path))
+    self.assertTrue(stdout_mock.getvalue().startswith(expected_prefix))
 
   def test_benchmark_stem(self):
-    pass
+    module = import_example('benchmark_stem')
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      path = os.path.join(DESC_DIR, 'collector', 'server-descriptors-2005-12-cropped.tar')
+      expected_prefix = EXPECTED_SERVER_DESC_BENCHMARK_PREFIX % path
+
+      module.measure_average_advertised_bandwidth(path)
+      self.assertTrue(stdout_mock.getvalue().startswith(expected_prefix))
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      path = os.path.join(DESC_DIR, 'collector', 'extra-infos-2019-04-cropped.tar')
+      expected_prefix = EXPECTED_EXTRAINFO_BENCHMARK_PREFIX % path
+
+      module.measure_countries_v3_requests(path)
+      self.assertTrue(stdout_mock.getvalue().startswith(expected_prefix))
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      path = os.path.join(DESC_DIR, 'collector', 'consensuses-2018-06-cropped.tar')
+      expected_prefix = EXPECTED_CONSENSUS_BENCHMARK_PREFIX % path
+
+      module.measure_average_relays_exit(path)
+      self.assertTrue(stdout_mock.getvalue().startswith(expected_prefix))
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      path = os.path.join(DESC_DIR, 'collector', 'microdescs-2019-05-cropped.tar')
+      expected_prefix = EXPECTED_MICRODESC_BENCHMARK_PREFIX % path
+
+      module.measure_fraction_relays_exit_80_microdescriptors(path)
+      self.assertTrue(stdout_mock.getvalue().startswith(expected_prefix))
 
   def test_broken_listener(self):
     pass