[tor-commits] [stem/master] Download method for the consensus

[atagar at torproject.org]

commit 69a2df402790653dadd26b301acc2bf1db96214e
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Aug 10 15:40:13 2019 -0700

    Download method for the consensus
---
 stem/descriptor/collector.py       | 61 +++++++++++++++++++++++++++++++++++++-
 test/integ/descriptor/collector.py | 31 +++++++++++++++++--
 2 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/stem/descriptor/collector.py b/stem/descriptor/collector.py
index 85eafec6..c0c97832 100644
--- a/stem/descriptor/collector.py
+++ b/stem/descriptor/collector.py
@@ -51,7 +51,8 @@ With this you can either download and read directly from CollecTor...
   get_instance - Provides a singleton CollecTor used for...
     |- get_server_descriptors - published server descriptors
     |- get_extrainfo_descriptors - published extrainfo descriptors
-    +- get_microdescriptors - published microdescriptors
+    |- get_microdescriptors - published microdescriptors
+    +- get_consensus - published router status entries
 
   File - Individual file residing within CollecTor
     |- read - provides descriptors from this file
@@ -61,6 +62,7 @@ With this you can either download and read directly from CollecTor...
     |- get_server_descriptors - published server descriptors
     |- get_extrainfo_descriptors - published extrainfo descriptors
     |- get_microdescriptors - published microdescriptors
+    |- get_consensus - published router status entries
     |
     |- index - metadata for content available from CollecTor
     +- files - files available from CollecTor
@@ -176,6 +178,17 @@ def get_microdescriptors(start = None, end = None, cache_to = None, timeout = No
     yield desc
 
 
+def get_consensus(start = None, end = None, cache_to = None, version = 3, microdescriptor = False, timeout = None, retries = 3):
+  """
+  Shorthand for
+  :func:`~stem.descriptor.collector.CollecTor.get_consensus`
+  on our singleton instance.
+  """
+
+  for desc in get_instance().get_consensus(start, end, cache_to, version, microdescriptor, timeout, retries):
+    yield desc
+
+
 class File(object):
   """
   File within CollecTor.
@@ -419,6 +432,8 @@ class CollecTor(object):
     :raises: :class:`~stem.DownloadFailed` if the download fails
     """
 
+    # TODO: support bridge variants ('bridge-server-descriptor' type)
+
     for f in self.files('server-descriptor', start, end):
       for desc in f.read(cache_to, timeout = timeout, retries = retries):
         yield desc
@@ -443,6 +458,8 @@ class CollecTor(object):
     :raises: :class:`~stem.DownloadFailed` if the download fails
     """
 
+    # TODO: support bridge variants ('bridge-extra-info' type)
+
     for f in self.files('extra-info', start, end):
       for desc in f.read(cache_to, timeout = timeout, retries = retries):
         yield desc
@@ -480,6 +497,48 @@ class CollecTor(object):
       for desc in f.read(cache_to, timeout = timeout, retries = retries):
         yield desc
 
+  def get_consensus(self, start = None, end = None, cache_to = None, version = 3, microdescriptor = False, timeout = None, retries = 3):
+    """
+    Provides consensus router status entries published during the given time
+    range, sorted oldest to newest.
+
+    :param datetime.datetime start: time range to begin with
+    :param datetime.datetime end: time range to end with
+    :param str cache_to: directory to cache archives into, if an archive is
+      available here it is not downloaded
+    :param int version: consensus variant to retrieve (versions 2 or 3)
+    :param bool microdescriptor: provides the microdescriptor consensus if
+      **True**, standard consensus otherwise
+    :param int timeout: timeout for downloading each individual archive when
+      the connection becomes idle, no timeout applied if **None**
+    :param int retires: maximum attempts to impose on a per-archive basis
+
+    :returns: **iterator** of
+      :class:`~stem.descriptor.router_status_entry.RouterStatusEntry`
+      for the given time range
+
+    :raises: :class:`~stem.DownloadFailed` if the download fails
+    """
+
+    if version == 3 and not microdescriptor:
+      desc_type = 'network-status-consensus-3'
+    elif version == 3 and microdescriptor:
+      desc_type = 'network-status-microdesc-consensus-3'
+    elif version == 2 and not microdescriptor:
+      desc_type = 'network-status-2'
+    else:
+      if microdescriptor and version != 3:
+        raise ValueError('Only v3 microdescriptors are available (not version %s)' % version)
+      else:
+        raise ValueError('Only v2 and v3 router status entries are available (not version %s)' % version)
+
+    # TODO: support bridge variants ('bridge-network-status' type)
+    # TODO: document vs router status entries (ie. DocumentType)?
+
+    for f in self.files(desc_type, start, end):
+      for desc in f.read(cache_to, timeout = timeout, retries = retries):
+        yield desc
+
   def index(self, compression = 'best'):
     """
     Provides the archives available in CollecTor.
diff --git a/test/integ/descriptor/collector.py b/test/integ/descriptor/collector.py
index d38ce7d1..f1336c53 100644
--- a/test/integ/descriptor/collector.py
+++ b/test/integ/descriptor/collector.py
@@ -3,6 +3,7 @@ Integration tests for stem.descriptor.collector.
 """
 
 import datetime
+import re
 import unittest
 
 import test.require
@@ -56,8 +57,34 @@ class TestCollector(unittest.TestCase):
   def test_downloading_microdescriptors(self):
     recent_descriptors = list(stem.descriptor.collector.get_microdescriptors(start = RECENT))
 
-    if not (300 < len(recent_descriptors) < 800):
-      self.fail('Downloaded %i descriptors, expected 300-800' % len(recent_descriptors))  # 23 on 8/7/19
+    if not (10 < len(recent_descriptors) < 100):
+      self.fail('Downloaded %i descriptors, expected 10-100' % len(recent_descriptors))  # 23 on 8/7/19
+
+  @test.require.only_run_once
+  @test.require.online
+  def test_downloading_consensus_v3(self):
+    recent_descriptors = list(stem.descriptor.collector.get_consensus(start = RECENT))
+
+    if not (3000 < len(recent_descriptors) < 10000):
+      self.fail('Downloaded %i descriptors, expected 3000-10000' % len(recent_descriptors))  # 6554 on 8/10/19
+
+  @test.require.only_run_once
+  @test.require.online
+  def test_downloading_consensus_micro(self):
+    recent_descriptors = list(stem.descriptor.collector.get_consensus(start = RECENT, microdescriptor = True))
+
+    if not (3000 < len(recent_descriptors) < 10000):
+      self.fail('Downloaded %i descriptors, expected 3000-10000' % len(recent_descriptors))
+
+  def test_downloading_consensus_invalid_type(self):
+    test_values = (
+      ({'version': 2, 'microdescriptor': True}, 'Only v3 microdescriptors are available (not version 2)'),
+      ({'version': 1}, 'Only v2 and v3 router status entries are available (not version 1)'),
+      ({'version': 4}, 'Only v2 and v3 router status entries are available (not version 4)'),
+    )
+
+    for args, expected_msg in test_values:
+      self.assertRaisesRegexp(ValueError, re.escape(expected_msg), list, stem.descriptor.collector.get_consensus(**args))
 
   def _test_index(self, compression):
     if compression and not compression.available: