[tor-commits] [stem/master] Adding a 'content' attribute to the Query class

[atagar at torproject.org]

commit e7c165ca7122a12fd5baade33221887176de3fe7
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Aug 4 12:04:22 2013 -0700

    Adding a 'content' attribute to the Query class
    
    Including a Query attribute so the caller can get our raw descriptor content.
    I'm actually doing this to simplify the run() method's usage. It had provided
    an iterator so...
    
      query = Query(my_resource)
      print list(query)  # this would print the ist of descriptors
      print list(query)  # this would be an empty list
    
    The reason the second call is an empty list is because we already iterated over
    the query. This is confusing, especially since...
    
      query = Query(my_resource)
      query.run(True)
      print list(query)  # also an empty list
    
    ... due to run() returning a list under the covers. By making run() provide a
    fresh iterator each time I also now had the downloaded content handy so making
    it public made sense (it'll commonly be wanted by our callers).
---
 stem/descriptor/remote.py      |   22 ++++++++++++----------
 test/unit/descriptor/remote.py |   16 ++++++++++++++++
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/stem/descriptor/remote.py b/stem/descriptor/remote.py
index 23a1f8c..95934e1 100644
--- a/stem/descriptor/remote.py
+++ b/stem/descriptor/remote.py
@@ -200,6 +200,7 @@ class Query(object):
   :var bool fall_back_to_authority: when retrying request issues the last
     request to a directory authority if **True**
 
+  :var str content: downloaded descriptor content
   :var Exception error: exception if a problem occured
   :var bool is_done: flag that indicates if our request has finished
   :var str download_url: last url used to download the descriptor, this is
@@ -235,6 +236,7 @@ class Query(object):
     self.retries = retries
     self.fall_back_to_authority = fall_back_to_authority
 
+    self.content = None
     self.error = None
     self.is_done = False
     self.download_url = None
@@ -249,8 +251,6 @@ class Query(object):
     self._downloader_thread = None
     self._downloader_thread_lock = threading.RLock()
 
-    self._results = None  # descriptor iterator
-
     if start:
       self.start()
 
@@ -307,14 +307,21 @@ class Query(object):
 
         raise self.error
       else:
-        if self._results is None:
+        if self.content is None:
           if suppress:
             return
 
           raise ValueError('BUG: _download_descriptors() finished without either results or an error')
 
         try:
-          for desc in self._results:
+          results = stem.descriptor.parse_file(
+            io.BytesIO(self.content),
+            self.descriptor_type,
+            validate = self.validate,
+            document_handler = self.document_handler,
+          )
+
+          for desc in results:
             yield desc
         except ValueError as exc:
           self.error = exc  # encountered a parsing error
@@ -357,12 +364,7 @@ class Query(object):
       if self.download_url.endswith('.z'):
         response = zlib.decompress(response)
 
-      self._results = stem.descriptor.parse_file(
-        io.BytesIO(response.strip()),
-        self.descriptor_type,
-        validate = self.validate,
-        document_handler = self.document_handler,
-      )
+      self.content = response.strip()
 
       self.runtime = time.time() - self.start_time
       log.trace("Descriptors retrieved from '%s' in %0.2fs" % (self.download_url, self.runtime))
diff --git a/test/unit/descriptor/remote.py b/test/unit/descriptor/remote.py
index 0d57dc3..e20dd06 100644
--- a/test/unit/descriptor/remote.py
+++ b/test/unit/descriptor/remote.py
@@ -123,3 +123,19 @@ class TestDescriptorDownloader(unittest.TestCase):
       timeout = 5,
     )
     self.assertEqual(3, urlopen_mock.call_count)
+
+  @patch('urllib2.urlopen')
+  def test_can_iterate_multiple_times(self, urlopen_mock):
+    urlopen_mock.return_value = io.BytesIO(TEST_DESCRIPTOR)
+
+    query = stem.descriptor.remote.Query(
+      '/tor/server/fp/9695DFC35FFEB861329B9F1AB04C46397020CE31',
+      'server-descriptor 1.0',
+      endpoints = [('128.31.0.39', 9131)],
+    )
+
+    # check that iterating over the query provides the descriptors each time
+
+    self.assertEqual(1, len(list(query)))
+    self.assertEqual(1, len(list(query)))
+    self.assertEqual(1, len(list(query)))