Fixed bug in csvutil.DictWriter; added tests.

rattail/csvutil.py

@@ -39,19 +39,15 @@ import codecs
 
 class DictWriter(csv.DictWriter):
     """
-    Convenience class which derives from the standard ``csv.DictWriter``.  This
-    currently exists only to provide the :meth:`writeheader()` method.
+    Convenience implementation of ``csv.DictWriter``.
+
+    This exists only to provide the :meth:`writeheader()` method on Python 2.6.
     """
 
     def writeheader(self):
-        """
-        Provide a ``writeheader()`` method in case the ``csv`` library in use
-        doesn't (i.e. for Python < 2.7).
-        """
-
         if hasattr(csv.DictWriter, 'writeheader'):
             return csv.DictWriter.writeheader(self)
-        self.writerow(self.fieldnames)
+        self.writer.writerow(self.fieldnames)
 
 
 class UTF8Recoder(object):

rattail/tests/__init__.py

@@ -1,19 +0,0 @@
-
-import unittest
-
-from pyramid import testing
-
-
-class TestCase(unittest.TestCase):
-    """
-    Base class for all test suites.
-    """
-
-    def setUp(self):
-        self.config = testing.setUp()
-
-    def tearDown(self):
-        testing.tearDown()
-
-    def test_something(self):
-        self.assertTrue(1)

tests/test_csvutil.py

@@ -0,0 +1,30 @@
+
+from unittest import TestCase
+from mock import patch
+
+from rattail import csvutil
+from cStringIO import StringIO
+
+
+class TestDictWriter(TestCase):
+
+    def test_writeheader_26(self):
+        # Simulate Python 2.6
+        with patch('csv.writer'):
+            with patch('rattail.csvutil.csv.DictWriter', spec=['writer']) as DictWriter:
+                buf = StringIO()
+                writer = csvutil.DictWriter(buf, ['field1', 'field2'])
+                writer.writeheader()
+                buf.close()
+                writer.writer.writerow.assert_called_once_with(['field1', 'field2'])
+
+    def test_writeheader_27(self):
+        # Simulate Python 2.7+
+        with patch('csv.writer'):
+            with patch('rattail.csvutil.csv.DictWriter', spec=['writer', 'writeheader']) as DictWriter:
+                buf = StringIO()
+                writer = csvutil.DictWriter(buf, ['field1', 'field2'])
+                writer.writeheader()
+                buf.close()
+                self.assertFalse(writer.writer.writerow.called)
+                DictWriter.writeheader.assert_called_once_with(writer)