Add pytest-mock and fix build errors.

airbnb · Aug 20, 2018 · 1092648 · 1092648
1 parent 4403745
commit 1092648
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/omniduct/databases/_cursor_formatters.py b/omniduct/databases/_cursor_formatters.py
@@ -1,6 +1,7 @@
 import csv
 import io
 import pickle
+import six
 from distutils.version import LooseVersion
 
 import pandas as pd
@@ -149,7 +150,7 @@ class CsvCursorFormatter(CursorFormatter):
     }
 
     def init(self, include_header=True):
-        self.output = io.StringIO()
+        self.output = io.StringIO() if six.PY3 else io.BytesIO()
         self.include_header = include_header
         self.writer = csv.writer(self.output, **self.FORMAT_PARAMS)
 

diff --git a/tox.ini b/tox.ini
@@ -11,6 +11,7 @@ deps=
     pyfakefs
     pytest
     pytest-cov
+    pytest-mock
     requests
 commands=
     pytest --cov omniduct --cov-report term-missing tests