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

[atagar at torproject.org]

commit c8b543368352a013c17cfe41baf6e370ee475826
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Sep 28 17:18:37 2020 -0700

    Test download_descriptor example
---
 docs/_static/example/download_descriptor.py | 12 +++++-------
 test/unit/examples.py                       | 30 ++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/docs/_static/example/download_descriptor.py b/docs/_static/example/download_descriptor.py
index 7e712ecf..89ea4a6c 100644
--- a/docs/_static/example/download_descriptor.py
+++ b/docs/_static/example/download_descriptor.py
@@ -57,9 +57,6 @@ def parse(argv):
 
   for opt, arg in recognized_args:
     if opt in ('-t', '--type'):
-      if arg not in VALID_TYPES:
-        raise ValueError("'%s' isn't a recognized decriptor type, options are: %s" % (arg, ', '.join(VALID_TYPES)))
-
       args['descriptor_type'] = arg
     elif opt in ('-f', '--fingerprint'):
       if not stem.util.tor_tools.is_valid_fingerprint(arg):
@@ -88,9 +85,9 @@ def parse(argv):
   return Args(**args)
 
 
-def main():
+def main(argv):
   try:
-    args = parse(sys.argv[1:])
+    args = parse(argv)
   except ValueError as exc:
     print(exc)
     sys.exit(1)
@@ -125,7 +122,8 @@ def main():
     print("'%s' is not a recognized descriptor type, options are: %s" % (args.descriptor_type, ', '.join(VALID_TYPES)))
     sys.exit(1)
 
-  print(desc)
+  if desc:
+    print(desc)
 
 if __name__ == '__main__':
-  main()
+  main(sys.argv[1:])
diff --git a/test/unit/examples.py b/test/unit/examples.py
index af049f57..eb9bdb2a 100644
--- a/test/unit/examples.py
+++ b/test/unit/examples.py
@@ -115,6 +115,18 @@ moria1 has the Running flag but maatuska doesn't: DCAEC3D069DC39AAE43D13C8AF31B5
 maatuska has the Running flag but moria1 doesn't: E2BB13AA2F6960CD93ABE5257A825687F3973C62
 """
 
+EXPECTED_DOWNLOAD_DESCRIPTOR_UNKNOWN_TYPE = """\
+Downloading kaboom descriptor from 128.31.0.34:9131...
+
+'kaboom' is not a recognized descriptor type, options are: server, extrainfo, consensus
+"""
+
+EXPECTED_DOWNLOAD_DESCRIPTOR_PREFIX = """\
+Downloading server descriptor from 1.2.3.4:443...
+
+router caerSidi 71.35.133.197 9001 0 0
+"""
+
 EXPECTED_EXIT_USED = """\
 Tracking requests for tor exits. Press 'enter' to end.
 
@@ -484,8 +496,24 @@ class TestExamples(unittest.TestCase):
 
     self.assertEqual('found relay caerSidi (A7569A83B5706AB1B1A9CB52EFF7D2D32E4553EB)\n', stdout_mock.getvalue())
 
+  @patch('sys.exit', Mock())
   def test_download_descriptor(self):
-    pass
+    import download_descriptor
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      download_descriptor.main(['--help'])
+      self.assertTrue(stdout_mock.getvalue().startswith("Downloads a descriptor through Tor's ORPort"))
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      download_descriptor.main(['--type', 'kaboom'])
+      self.assertEqual(EXPECTED_DOWNLOAD_DESCRIPTOR_UNKNOWN_TYPE, stdout_mock.getvalue())
+
+    server_desc = RelayDescriptor.create({'router': 'caerSidi 71.35.133.197 9001 0 0'})
+
+    with patch('sys.stdout', new_callable = io.StringIO) as stdout_mock:
+      with patch('stem.descriptor.remote.get_server_descriptors', _download_of(server_desc)):
+        download_descriptor.main(['--dirport', '1.2.3.4:443'])
+        self.assertTrue(stdout_mock.getvalue().startswith(EXPECTED_DOWNLOAD_DESCRIPTOR_PREFIX))
 
   def test_ephemeral_hidden_services(self):
     pass