Refactor Python 2/3 compatible test code

Add compat.bom_prefix_csv to add a BOM to CSV text
Add compat.force_utf8 to optionally encode text in Py3

multigtfs/compat.py

@@ -2,10 +2,11 @@
 
 Handle compatibility between Python versions, Django versions, etc.
 """
+from codecs import BOM_UTF8
 from distutils.version import LooseVersion
 
 from django import get_version
-from django.utils.six import PY3
+from django.utils.six import PY3, binary_type
 
 DJ_VERSION = LooseVersion(get_version())
 
@@ -35,6 +36,29 @@ def _get_blank_value_19(field):
     get_blank_value = _get_blank_value_18
 
 
+def bom_prefix_csv(text):
+    """
+    Prefix CSV text with a Byte-order Marker (BOM).
+
+    The return value needs to be encoded differently so the CSV reader will
+    handle the BOM correctly:
+    - Python 2 returns a UTF-8 encoded bytestring
+    - Python 3 returns unicode text
+    """
+    if PY3:
+        return BOM_UTF8.decode('utf-8') + text
+    else:
+        return BOM_UTF8 + text.encode('utf-8')
+
+
+def force_utf8(text):
+    """Encode as UTF-8 bytestring if it isn't already."""
+    if isinstance(text, binary_type):
+        return text
+    else:
+        return text.encode('utf-8')
+
+
 def opener_from_zipfile(zipfile):
     """
     Returns a function that will open a file in a zipfile by name.

multigtfs/tests/test_agency.py

@@ -14,12 +14,12 @@
 # limitations under the License.
 
 from __future__ import unicode_literals
-from codecs import BOM_UTF8
 
 from django.test import TestCase
-from django.utils.six import StringIO, PY3
+from django.utils.six import StringIO
 
 from multigtfs.models import Agency, Feed
+from multigtfs.compat import bom_prefix_csv
 
 
 class AgencyTest(TestCase):
@@ -86,17 +86,10 @@ def test_import_agency_maximal(self):
         self.assertEqual(agency.fare_url, 'http://google.com')
 
     def test_import_bom(self):
-        if PY3:  # pragma: no cover
-            text = (BOM_UTF8.decode('utf-8') + """\
+        agency_txt = StringIO(bom_prefix_csv("""\
 agency_name,agency_url,agency_timezone
 Demo Transit Authority,http://google.com,America/Los_Angeles
-""")
-        else:
-            text = (BOM_UTF8 + b"""\
-agency_name,agency_url,agency_timezone
-Demo Transit Authority,http://google.com,America/Los_Angeles
-""")
-        agency_txt = StringIO(text)
+"""))
         Agency.import_txt(agency_txt, self.feed)
         agency = Agency.objects.get()
         self.assertEqual(agency.agency_id, '')

multigtfs/tests/test_stop.py

@@ -14,12 +14,12 @@
 # limitations under the License.
 
 from __future__ import unicode_literals
-from codecs import BOM_UTF8
 
 from django.contrib.gis.geos import MultiLineString
 from django.test import TestCase
-from django.utils.six import StringIO, text_type, PY3
+from django.utils.six import StringIO
 
+from multigtfs.compat import bom_prefix_csv, force_utf8
 from multigtfs.models import Feed, Route, Stop, StopTime, Trip, Zone
 
 
@@ -200,25 +200,14 @@ def test_import_stops_txt_stop_before_station_plus_extra(self):
         self.assertEqual(expected, self.feed.meta)
 
     def test_import_stops_txt_bom(self):
-        if PY3:  # pragma: no cover
-            text = (BOM_UTF8.decode('utf-8') + """\
+        stops_txt = StringIO(bom_prefix_csv("""\
 stop_id,stop_code,stop_name,stop_desc,stop_lat,stop_lon,zone_id,stop_url,\
 location_type,parent_station,stop_timezone
 FUR_CREEK_RES,FC,Furnace Creek Resort,,36.425288,-117.133162,A,\
 http://example.com/fcr,0,FUR_CREEK_STA,
 FUR_CREEK_STA,,Furnace Creek Station,"Our Station",36.425288,-117.133162,A,\
 http://example.com,1,,America/Los_Angeles
-""")
-        else:
-            text = (BOM_UTF8 + b"""\
-stop_id,stop_code,stop_name,stop_desc,stop_lat,stop_lon,zone_id,stop_url,\
-location_type,parent_station,stop_timezone
-FUR_CREEK_RES,FC,Furnace Creek Resort,,36.425288,-117.133162,A,\
-http://example.com/fcr,0,FUR_CREEK_STA,
-FUR_CREEK_STA,,Furnace Creek Station,"Our Station",36.425288,-117.133162,A,\
-http://example.com,1,,America/Los_Angeles
-""")
-        stops_txt = StringIO(text)
+"""))
         Stop.import_txt(stops_txt, self.feed)
         self.assertEqual(Stop.objects.count(), 2)
         station = Stop.objects.get(stop_id='FUR_CREEK_STA')
@@ -258,9 +247,7 @@ def test_export_stops_utf8(self):
             feed=self.feed, stop_id=6071,
             name=b'The Delta Caf\x82'.decode('latin1'),
             point='POINT(-95.975834 36.114554)')
-        stops_txt = Stop.export_txt(self.feed)
-        if isinstance(stops_txt, text_type):  # pragma: no cover
-            stops_txt = stops_txt.encode('utf-8')
+        stops_txt = force_utf8(Stop.export_txt(self.feed))
         self.assertEqual(stops_txt, b"""\
 stop_id,stop_name,stop_lat,stop_lon
 6071,The Delta Caf\xc2\x82,36.114554,-95.975834