commit 7022021c4207a0066cb93261629f2ba020d307f6
Author: Damian Johnson <atagar(a)torproject.org>
Date: Sat Aug 4 12:13:47 2012 -0700
Revising descriptor exporter and associated tests
Several changes for the prior csv export functionality, especially its tests.
Most notably this includes...
* rewrite of the tests to be simpler and break down by test cases
* newline breaks rather than windows style '\r\n'
* excluding private attributes (attributes that start with an underscore)
* writing directly to a file in export_csv_file() rather than buffering to a string first
* general refactoring to hopefully improve readability
---
run_tests.py | 6 +-
stem/descriptor/__init__.py | 2 +-
stem/descriptor/export.py | 148 ++++++++++++++----------------
test/unit/descriptor/__init__.py | 2 +-
test/unit/descriptor/export.py | 189 ++++++++++++++------------------------
5 files changed, 142 insertions(+), 205 deletions(-)
diff --git a/run_tests.py b/run_tests.py
index bed8e50..4ad551a 100755
--- a/run_tests.py
+++ b/run_tests.py
@@ -16,6 +16,7 @@ import test.output
import test.runner
import test.check_whitespace
import test.unit.connection.authentication
+import test.unit.descriptor.export
import test.unit.descriptor.reader
import test.unit.descriptor.server_descriptor
import test.unit.descriptor.extrainfo_descriptor
@@ -50,8 +51,6 @@ import test.integ.util.proc
import test.integ.util.system
import test.integ.process
import test.integ.version
-#TODO move this by other descriptor tests
-import test.unit.descriptor.export
import stem.prereq
import stem.util.conf
@@ -109,6 +108,7 @@ UNIT_TESTS = (
test.unit.util.proc.TestProc,
test.unit.util.system.TestSystem,
test.unit.util.tor_tools.TestTorTools,
+ test.unit.descriptor.export.TestExport,
test.unit.descriptor.reader.TestDescriptorReader,
test.unit.descriptor.server_descriptor.TestServerDescriptor,
test.unit.descriptor.extrainfo_descriptor.TestExtraInfoDescriptor,
@@ -123,9 +123,7 @@ UNIT_TESTS = (
test.unit.response.protocolinfo.TestProtocolInfoResponse,
test.unit.response.authchallenge.TestAuthChallengeResponse,
test.unit.connection.authentication.TestAuthenticate,
- test.unit.descriptor.export.TestExport,
)
-# TODO move export test to proper location
INTEG_TESTS = (
test.integ.util.conf.TestConf,
diff --git a/stem/descriptor/__init__.py b/stem/descriptor/__init__.py
index bb6f874..45cd016 100644
--- a/stem/descriptor/__init__.py
+++ b/stem/descriptor/__init__.py
@@ -13,7 +13,7 @@ Package for parsing and processing descriptor data.
"""
__all__ = [
- "descriptor",
+ "export",
"reader",
"extrainfo_descriptor",
"server_descriptor",
diff --git a/stem/descriptor/export.py b/stem/descriptor/export.py
index 89a5644..16a148d 100644
--- a/stem/descriptor/export.py
+++ b/stem/descriptor/export.py
@@ -1,97 +1,89 @@
-import os, csv, sets, cStringIO
+"""
+Utilities for exporting descriptors to other formats.
+"""
+import csv
+import cStringIO
-def export_csv(descriptor, include_fields=(), exclude_fields=()):
+import stem.descriptor
+
+class _ExportDialect(csv.excel):
+ lineterminator = '\n'
+
+def export_csv(descriptors, included_fields = (), excluded_fields = (), header = True):
"""
- Takes a single descriptor object, puts it in a list, and passes it to
- descriptors_csv_exp to build a csv.
+ Provides a newline separated CSV for one or more descriptors. If simply
+ provided with descriptors then the CSV contains all of its attributes,
+ labelled with a header row. Either 'included_fields' or 'excluded_fields' can
+ be used for more granular control over its attributes and the order.
- :param object descriptor: single descriptor whose attributes will be returned as a string.
- :param list include_fields: list of attribute fields to include in the csv string.
- :param list exclude_fields: list of attribute fields to exclude from csv string.
+ :param Descriptor,list descriptors: descriptor or list of descriptors to be exported
+ :param list included_fields: attributes to include in the csv
+ :param list excluded_fields: attributes to exclude from the csv
+ :param bool header: if True then the first line will be a comma separated list of the attribute names
- :returns: single csv line as a string with one descriptor attribute per cell.
+ :returns: str of the CSV for the descriptors, one per line
+ :raises: ValueError if descriptors contain more than one descriptor type
"""
- descr = (descriptor,)
- return export_csvs(descr, include_fields=include_fields, exclude_fields=exclude_fields)
+ output_buffer = cStringIO.StringIO()
+ export_csv_file(output_buffer, descriptors, included_fields, excluded_fields, header)
+ return output_buffer.getvalue()
-
-def export_csvs(descriptors, include_fields=[], exclude_fields=[], header=False):
+def export_csv_file(output_file, descriptors, included_fields = (), excluded_fields = (), header = True):
"""
- Takes an iterable of descriptors, returns a string with one line per descriptor
- where each line is a comma separated list of descriptor attributes.
+ Similar to :func:`stem.descriptor.export.export_csv`, except that the CSV is
+ written directly to a file.
- :param list descrs: List of descriptor objects whose attributes will be written.
- :param list include_fields: list of attribute fields to include in the csv string.
- :param list exclude_fields: list of attribute fields to exclude from csv string.
- :param bool header: whether or not a header is requested.
+ :param file output_file: file to be written to
+ :param Descriptor,list descriptors: descriptor or list of descriptors to be exported
+ :param list included_fields: attributes to include in the csv
+ :param list excluded_fields: attributes to exclude from the csv
+ :param bool header: if True then the first line will be a comma separated list of the attribute names
- :returns: csv string with one descriptor per line and one attribute per cell.
- :raises: ValueError if more than one descriptor type (e.g. server_descriptor,
- extrainfo_descriptor) is provided in the iterable.
+ :returns: str of the CSV for the descriptors, one per line
+ :raises: ValueError if descriptors contain more than one descriptor type
"""
- # Need a file object to write to with DictWriter.
- temp_file = cStringIO.StringIO()
+ if isinstance(descriptors, stem.descriptor.Descriptor):
+ descriptors = (descriptors,)
- first = True
+ if not descriptors:
+ return
- for desc in descriptors:
- #umport sys
- attr = vars(desc)
-
- # Defining incl_fields and the dwriter object requires having access
- # to a descriptor object.
- if first:
- # All descriptor objects should be of the same type
- # (i.e. server_descriptor.RelayDesrciptor)
- desc_type = type(desc)
-
- # define incl_fields, 4 cases where final case is incl_fields already
- # defined and excl_fields left blank, so no action is necessary.
- if not include_fields and exclude_fields:
- incl = set(attr.keys())
- include_fields = list(incl.difference(exclude_fields))
-
- elif not include_fields and not exclude_fields:
- include_fields = attr.keys()
-
- elif include_fields and exclude_fields:
- incl = set(include_fields)
- include_fields = list(incl.difference(exclude_fields))
-
- dwriter = csv.DictWriter(temp_file, include_fields, extrasaction='ignore')
-
- if header:
- dwriter.writeheader()
- first = False
-
- if desc_type == type(desc):
- dwriter.writerow(attr)
- else:
- raise ValueError('More than one descriptor type provided. Started with a %s, and just got a %s' % (desc_type, type(desc)))
+ descriptor_type = type(descriptors[0])
+ descriptor_type_label = descriptor_type.__name__
+ included_fields = list(included_fields)
- return temp_file.getvalue()
- # cStringIO files are closed automatically when the current scope is exited.
-
-def export_csv_file(descriptors, document, include_fields=(), exclude_fields=(), header=True):
- """
- Writes descriptor attributes to a csv file on disk.
+ # If the user didn't specify the fields to include then export everything,
+ # ordered alphabetically. If they did specify fields then make sure that
+ # they exist.
- Calls get_csv_lines with the given argument, then writes the returned string
- to a file location specified by document_location.
- Precondition that param document has a 'write' attribute.
+ desc_attr = sorted(vars(descriptors[0]).keys())
- :param list descriptors: descriptor objects with attributes to export as csv file.
- :param object document: File object to be written to.
- :param bool header: defaults to true, determines if document will have a header row.
- :param list include_fields: list of attribute fields to include in the csv line.
- :param list exclude_fields: list of attribute fields to exclude from csv line.
- """
+ if included_fields:
+ for field in included_fields:
+ if not field in desc_attr:
+ raise ValueError("%s does not have a '%s' attribute, valid fields are: %s" % (descriptor_type_label, field, ", ".join(desc_attr)))
+ else:
+ included_fields = [attr for attr in desc_attr if not attr.startswith('_')]
+
+ for field in excluded_fields:
+ try:
+ included_fields.remove(field)
+ except ValueError:
+ pass
+
+ writer = csv.DictWriter(output_file, included_fields, dialect = _ExportDialect(), extrasaction='ignore')
+
+ if header:
+ writer.writeheader()
- try:
- document.write(export_csvs(descriptors, include_fields=include_fields, exclude_fields=exclude_fields, header=header))
- except AttributeError:
- print "Provided %r object does not have a write() method." % document
- raise
+ for desc in descriptors:
+ if not isinstance(desc, stem.descriptor.Descriptor):
+ raise ValueError("Unable to export a descriptor CSV since %s is not a descriptor." % type(desc).__name__)
+ elif descriptor_type != type(desc):
+ raise ValueError("To export a descriptor CSV all of the descriptors must be of the same type. First descriptor was a %s but we later got a %s." % (descriptor_type_label, type(desc)))
+
+ writer.writerow(vars(desc))
+
diff --git a/test/unit/descriptor/__init__.py b/test/unit/descriptor/__init__.py
index cf4fceb..1837a07 100644
--- a/test/unit/descriptor/__init__.py
+++ b/test/unit/descriptor/__init__.py
@@ -2,5 +2,5 @@
Unit tests for stem.descriptor.
"""
-__all__ = ["reader", "extrainfo_descriptor", "server_descriptor"]
+__all__ = ["export", "reader", "extrainfo_descriptor", "server_descriptor"]
diff --git a/test/unit/descriptor/export.py b/test/unit/descriptor/export.py
index 12d5ed3..6ef52cb 100644
--- a/test/unit/descriptor/export.py
+++ b/test/unit/descriptor/export.py
@@ -1,146 +1,93 @@
"""
-Unit testing code for the stem.descriptor.export module
+Unit tests for stem.descriptor.export.
"""
-import unittest
-import stem.descriptor.export as export
-import test.mocking as mocking
-
-from collections import namedtuple
-import stem.descriptor as desc
-import cStringIO
-
-# Create descriptor objects.
-DESCR_DICT = {'average_bandwidth': 5242880, 'onion_key': 'RSA PUB = JAIK', 'address': '79.139.135.90', '_digest': None, 'exit_policy': ['reject *:*'], 'fingerprint': 'AAAAAAAAAAAAAAAAAAA'}
-DESCR2_DICT = {'average_bandwidth': 5555555, 'onion_key': 'RSA PUB = GOUS', 'address': '100.1.1.1', '_digest': None, 'exit_policy': ['reject *:*'], 'fingerprint': 'BBBBBBBBBBBBBBBBBBBBB'}
-DESCR3_DICT = {'bandwidth':12345,'average_bandwidth': 6666666, 'address': '101.0.0.1','extra_info':None}
-RAW = 'router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon'
-
-descriptor = desc.Descriptor(RAW)
-descriptor.__dict__.update(DESCR_DICT)
-
-descriptor2 = desc.Descriptor(RAW)
-descriptor2.__dict__.update(DESCR2_DICT)
-
-descriptor3 = desc.server_descriptor.RelayDescriptor(RAW, validate=False)
-descriptor3.__dict__.update(DESCR3_DICT)
-
-# Expected return csv strings.
-SINGLE_ALL = '5242880,RSA PUB = JAIK,AAAAAAAAAAAAAAAAAAA,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],79.139.135.90,'
-SINGLE_PART = '79.139.135.90,[\'reject *:*\']'
-SINGLE_PART2 = '5242880,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],79.139.135.90,'
-SINGLE_PART3 = '79.139.135.90,AAAAAAAAAAAAAAAAAAA'
-DOUBLE_ALL = '5242880,RSA PUB = JAIK,AAAAAAAAAAAAAAAAAAA,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],79.139.135.90,\r\n5555555,RSA PUB = GOUS,BBBBBBBBBBBBBBBBBBBBB,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],100.1.1.1,\r\n'
-DOUBLE_PART = '79.139.135.90,[\'reject *:*\']\r\n100.1.1.1,[\'reject *:*\']\r\n'
-DOUBLE_PART2 = '5242880,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],79.139.135.90,\r\n5555555,,router TORsinn3r 46.17.96.217 9001 0 0 platform Tor 0.2.3.19-rc on Linux bandwidth 4 5 6 ...andonandon,[\'reject *:*\'],100.1.1.1,\r\n'
+import StringIO
+import unittest
-SINGLE_ALL_HEAD = 'average_bandwidth,onion_key,fingerprint,_digest,_raw_contents,exit_policy,address,_path\r\n' + SINGLE_ALL + '\r\n'
-SINGLE_PART3_HEAD = 'address,fingerprint\r\n' + SINGLE_PART3
-DOUBLE_ALL_HEAD = 'average_bandwidth,onion_key,fingerprint,_digest,_raw_contents,exit_policy,address,_path\r\n' + DOUBLE_ALL
-DOUBLE_PART_HEAD = 'address,exit_policy\r\n' + DOUBLE_PART
-DOUBLE_PART2_HEAD = 'average_bandwidth,_digest,_raw_contents,exit_policy,address,_path\r\n' + DOUBLE_PART2
+from stem.descriptor.export import export_csv, export_csv_file
+from stem.descriptor.server_descriptor import RelayDescriptor, BridgeDescriptor
+from test.mocking import get_server_descriptor
class TestExport(unittest.TestCase):
- def tearDown(self):
- mocking.revert_mocking()
-
- def test_export_csv(self):
+ def test_minimal_descriptor(self):
"""
- Tests the export_csv function which takes a single descriptor object.
+ Exports a single minimal tor server descriptor.
"""
- Fields = namedtuple('Fields', 'include_fields exclude_fields')
- # Descriptors must be an iterable
- # named tuples replace dictionaries as dict keys must immutable.
- ret_vals = {((descriptor,), Fields(include_fields=(), exclude_fields=())):SINGLE_ALL,
- ((descriptor,), Fields(include_fields=('address', 'exit_policy'),
- exclude_fields=())):SINGLE_PART,
- ((descriptor,), Fields(include_fields=(),
- exclude_fields=('onion_key', 'fingerprint'))):SINGLE_PART2,
- ((descriptor,), Fields(include_fields=('address', 'exit_policy', 'fingerprint'),
- exclude_fields=('fingerprint',))):SINGLE_PART,
- ((descriptor,), Fields(include_fields=('address', 'fingerprint'),
- exclude_fields=('_digest',))):SINGLE_PART3
- }
- mocking.mock(export.export_csvs, mocking.return_for_args(ret_vals, kwarg_type=Fields))
+ desc = RelayDescriptor(get_server_descriptor())
- # Used tuples for incl/exclude_fields for parameter matching with ret_vals dict.
- self.assertEqual(SINGLE_ALL, export.export_csv(descriptor))
- self.assertEqual(SINGLE_PART, export.export_csv(descriptor,
- include_fields=('address', 'exit_policy')))
- self.assertEqual(SINGLE_PART2, export.export_csv(descriptor,
- exclude_fields=('onion_key', 'fingerprint')))
- self.assertEqual(SINGLE_PART, export.export_csv(descriptor,
- include_fields=('address', 'exit_policy', 'fingerprint'), exclude_fields=('fingerprint',)))
- self.assertEqual(SINGLE_PART3, export.export_csv(descriptor,
- include_fields=('address', 'fingerprint'), exclude_fields=('_digest',)))
-
+ desc_csv = export_csv(desc, included_fields = ('nickname', 'address', 'published'), header = False)
+ expected = "caerSidi,71.35.133.197,2012-03-01 17:15:27\n"
+ self.assertEquals(expected, desc_csv)
+
+ desc_csv = export_csv(desc, included_fields = ('nickname', 'address', 'published'), header = True)
+ expected = "nickname,address,published\n" + expected
+ self.assertEquals(expected, desc_csv)
- def test_export_csvs(self):
+ def test_multiple_descriptors(self):
"""
- Test the export_csvs function which takes a list of descriptor objects.
+ Exports multiple descriptors, making sure that we get them back in the same
+ order.
"""
- # Single descriptor
- self.assertEquals(SINGLE_ALL + "\r\n", export.export_csvs([descriptor]))
- self.assertEqual(SINGLE_PART + "\r\n", export.export_csvs([descriptor],
- include_fields=['address', 'exit_policy']))
- self.assertEqual(SINGLE_PART2 + "\r\n", export.export_csvs([descriptor],
- exclude_fields=['onion_key', 'fingerprint']))
- self.assertEqual(SINGLE_PART + "\r\n", export.export_csvs([descriptor],
- include_fields=['address', 'exit_policy', 'fingerprint'], exclude_fields=['fingerprint']))
-
- # Multiple descriptors
- self.assertEqual(DOUBLE_ALL, export.export_csvs([descriptor, descriptor2]))
- self.assertEqual(DOUBLE_PART, export.export_csvs([descriptor, descriptor2],
- include_fields=['address', 'exit_policy']))
- self.assertEqual(DOUBLE_PART2, export.export_csvs([descriptor, descriptor2],
- exclude_fields=['onion_key', 'fingerprint']))
- self.assertEqual(DOUBLE_PART, export.export_csvs([descriptor, descriptor2],
- include_fields=['address', 'exit_policy', 'fingerprint'], exclude_fields=['fingerprint']))
+ nicknames = ('relay1', 'relay3', 'relay2', 'caerSidi', 'zeus')
+ descriptors = []
- # Tests with headers
- self.assertEqual(SINGLE_ALL_HEAD, export.export_csvs([descriptor], header=True))
- self.assertEqual(SINGLE_PART3_HEAD + "\r\n", export.export_csvs([descriptor],
- include_fields=['address', 'fingerprint'], exclude_fields=['_digest'], header=True))
- self.assertEqual(DOUBLE_ALL_HEAD, export.export_csvs([descriptor, descriptor2], header=True))
- self.assertEqual(DOUBLE_PART_HEAD, export.export_csvs([descriptor, descriptor2],
- include_fields=['address', 'exit_policy'], header=True))
- self.assertEqual(DOUBLE_PART2_HEAD, export.export_csvs([descriptor, descriptor2],
- exclude_fields=['onion_key', 'fingerprint'], header=True))
+ for nickname in nicknames:
+ router_line = "%s 71.35.133.197 9001 0 0" % nickname
+ descriptors.append(RelayDescriptor(get_server_descriptor({'router': router_line})))
- # Other tests
- self.assertRaises(ValueError, export.export_csvs, [descriptor, descriptor3])
- self.assertRaises(ValueError, export.export_csvs, [descriptor, descriptor3],
- include_fields=['onion_key', 'address', 'fingerprint'], exclude_fields=['onion_key'])
-
+ expected = "\n".join(nicknames) + "\n"
+ self.assertEqual(expected, export_csv(descriptors, included_fields = ('nickname',), header = False))
-
- def test_export_csv_file(self):
+ def test_file_output(self):
"""
- Tests the export_csv_file function.
+ Basic test for the export_csv_file() function, checking that it provides
+ the same output as export_csv().
"""
- sample_csv_string = 'This, is, a, sample, string.\r\nline, two.\r\n'
- sample_csv_string2 = 'Another, sample\r\n,, second,\r\n'
- sample_file = cStringIO.StringIO()
- # Must use named tuples again for ret_vals dictionary.
- Fields = namedtuple('Fields', 'include_fields exclude_fields header')
+ desc = RelayDescriptor(get_server_descriptor())
+ desc_csv = export_csv(desc)
- ret_vals = {((descriptor,), Fields(include_fields=(), exclude_fields=(), header=True)):sample_csv_string,
- ((descriptor,), Fields(include_fields=('address', 'onion_key'), exclude_fields=('address',), header=False)):sample_csv_string2}
- # TODO Ask Danner: mock it once then do both tests (not including assertRaises), or do separate mockings.
- # the latter requires that we still include empty incl_fields and excl_fields parameters instead of
- # letting them default to []. Same for header.
- mocking.mock(export.export_csvs, mocking.return_for_args(ret_vals, kwarg_type=Fields))
+ csv_buffer = StringIO.StringIO()
+ export_csv_file(csv_buffer, desc)
- export.export_csv_file((descriptor,), sample_file)
- self.assertEqual(sample_csv_string, sample_file.getvalue())
+ self.assertEqual(desc_csv, csv_buffer.getvalue())
+
+ def test_excludes_private_attr(self):
+ """
+ Checks that the default attributes for our csv output doesn't include private fields.
+ """
- sample_file = cStringIO.StringIO()
+ desc = RelayDescriptor(get_server_descriptor())
+ desc_csv = export_csv(desc)
- export.export_csv_file((descriptor,), sample_file, include_fields=('address', 'onion_key'), exclude_fields=('address',), header=False)
- self.assertEqual(sample_csv_string2, sample_file.getvalue())
+ self.assertTrue(',signature' in desc_csv)
+ self.assertFalse(',_digest' in desc_csv)
+ self.assertFalse(',_annotation_lines' in desc_csv)
+
+ def test_empty_input(self):
+ """
+ Exercises when we don't provide any descriptors.
+ """
- # Make sure error is Raised when necessary.
- self.assertRaises(AttributeError, export.export_csv_file, (descriptor,), sample_csv_string)
+ self.assertEquals("", export_csv([]))
+
+ def test_invalid_attributes(self):
+ """
+ Attempts to make a csv with attributes that don't exist.
+ """
+
+ desc = RelayDescriptor(get_server_descriptor())
+ self.assertRaises(ValueError, export_csv, desc, ('nickname', 'blarg!'))
+
+ def test_multiple_descriptor_types(self):
+ """
+ Attempts to make a csv with multiple descriptor types.
+ """
+
+ server_desc = RelayDescriptor(get_server_descriptor())
+ bridge_desc = BridgeDescriptor(get_server_descriptor(is_bridge = True))
+ self.assertRaises(ValueError, export_csv, (server_desc, bridge_desc))
+