Make unittests work with py3 (#140)

* Make some unittest work with py3

* Make some unit tests work with py3

* Make snowflake metadata extractor py3 test green

.travis.yml

 language: python
-python: 2.7
+python: 
+- '2.7'
+- '3.6'
 install:
 - pip install -r requirements.txt
 - python setup.py install

Makefile

@@ -4,7 +4,7 @@ clean:
 	rm -rf dist/
 
 .PHONY: test_unit
-test_unit: test_unit2
+test_unit: test_unit2_or3_if_its_default
 
 lint:
 	flake8 .
@@ -13,8 +13,8 @@ lint:
 test: test_unit lint
 
 .PHONY: test_unit
-test_unit2:
-	python -m pytest tests/unit
+test_unit2_or3_if_its_default:
+	python -bb -m pytest tests/unit
 
 test_unit3:
 	python3 -bb -m pytest tests/unit

databuilder/extractor/snowflake_metadata_extractor.py

 import logging
+import six
 from collections import namedtuple
 
 from pyhocon import ConfigFactory, ConfigTree  # noqa: F401
@@ -72,7 +73,9 @@ class SnowflakeMetadataExtractor(Extractor):
         else:
             cluster_source = "'{}'".format(self._cluster)
 
-        self._database = conf.get_string(SnowflakeMetadataExtractor.DATABASE_KEY).encode('utf-8', 'ignore')
+        self._database = conf.get_string(SnowflakeMetadataExtractor.DATABASE_KEY)
+        if six.PY2:
+            self._database = self._database.encode('utf-8', 'ignore')
 
         self.sql_stmt = SnowflakeMetadataExtractor.SQL_STATEMENT.format(
             where_clause_suffix=conf.get_string(SnowflakeMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY),

databuilder/models/user.py

@@ -131,7 +131,7 @@ class User(Neo4jCsvSerializable):
         result_node[User.USER_NODE_UPDATED_AT] = self.updated_at if self.updated_at else 0
 
         if self.attrs:
-            print (self.attrs)
+            print(self.attrs)
             for k, v in self.attrs.items():
                 if k not in result_node:
                     result_node[k] = v

setup.cfg

@@ -9,7 +9,7 @@ ignore = NONE
 max-line-length = 120
 
 [tool:pytest]
-addopts = --cov=databuilder --cov-fail-under=70 --cov-report=term-missing:skip-covered --cov-report=xml --cov-report=html -vvv
+addopts = -rs --cov=databuilder --cov-fail-under=70 --cov-report=term-missing:skip-covered --cov-report=xml --cov-report=html -vvv
 
 [coverage:run]
 branch = True

tests/unit/extractor/test_bigquery_usage_extractor.py

@@ -2,6 +2,8 @@ from mock import patch, Mock
 import base64
 import tempfile
 import unittest
+import six
+import pytest
 
 from pyhocon import ConfigFactory
 
@@ -180,6 +182,10 @@ class MockLoggingClient():
         return self.b
 
 
+@pytest.mark.skipif(
+    six.PY3,
+    reason="Deactivated for Python3 because of issue [#40](https://github.com/lyft/amundsen/issues/40) (ANTLR for py3)"
+)
 class TestBigqueryUsageExtractor(unittest.TestCase):
 
     @patch('databuilder.extractor.bigquery_usage_extractor.build')

tests/unit/extractor/test_table_column_usage_aggregate_extractor.py

@@ -2,13 +2,20 @@ import unittest
 
 from mock import patch, MagicMock  # noqa: F401
 from pyhocon import ConfigFactory
+import pytest
+import six
 
-from databuilder.extractor.table_column_usage_aggregate_extractor import TblColUsgAggExtractor, RAW_EXTRACTOR
-from databuilder.models.table_column_usage import TableColumnUsage, ColumnReader
-from databuilder.transformer.regex_str_replace_transformer import RegexStrReplaceTransformer
-from databuilder.transformer.sql_to_table_col_usage_transformer import SqlToTblColUsageTransformer
+if six.PY2:
+    from databuilder.extractor.table_column_usage_aggregate_extractor import TblColUsgAggExtractor, RAW_EXTRACTOR
+    from databuilder.models.table_column_usage import TableColumnUsage, ColumnReader
+    from databuilder.transformer.regex_str_replace_transformer import RegexStrReplaceTransformer
+    from databuilder.transformer.sql_to_table_col_usage_transformer import SqlToTblColUsageTransformer
 
 
+@pytest.mark.skipif(
+    six.PY3,
+    reason="Deactivated for Python3 because of issue [#40](https://github.com/lyft/amundsen/issues/40) (ANTLR for py3)"
+)
 class TestTblColUsgAggExtractor(unittest.TestCase):
 
     def test_aggregate(self):

tests/unit/loader/test_file_system_elasticsearch_json_loader.py

@@ -76,7 +76,7 @@ class TestFSElasticsearchJSONLoader(unittest.TestCase):
 
         with self.assertRaises(Exception) as context:
             loader.load(data)  # type: ignore
-        self.assertTrue("Record not of type 'ElasticsearchDocument'!" in context.exception)
+        self.assertIn("Record not of type 'ElasticsearchDocument'!", str(context.exception))
 
         loader.close()
 

tests/unit/publisher/test_elasticsearch_publisher.py

 import json
 from mock import MagicMock, mock_open, patch
+import six
 import unittest
 
 from pyhocon import ConfigFactory
@@ -34,7 +35,10 @@ class TestElasticsearchPublisher(unittest.TestCase):
         """
         Test Publish functionality with no data
         """
-        with patch('__builtin__.open', mock_open(read_data='')) as mock_file:
+        target = 'builtins.open'
+        if six.PY2:
+            target = '__builtin__.open'
+        with patch(target, mock_open(read_data='')) as mock_file:
             publisher = ElasticsearchPublisher()
             publisher.init(conf=Scoped.get_scoped_conf(conf=self.conf,
                                                        scope=publisher.get_scope()))
@@ -55,7 +59,10 @@ class TestElasticsearchPublisher(unittest.TestCase):
                                 'KEY_DOESNOT_MATTER2': 'NO_VALUE2'})
         self.mock_es_client.indices.get_alias.return_value = {}
 
-        with patch('__builtin__.open', mock_open(read_data=mock_data)) as mock_file:
+        target = 'builtins.open'
+        if six.PY2:
+            target = '__builtin__.open'
+        with patch(target, mock_open(read_data=mock_data)) as mock_file:
             publisher = ElasticsearchPublisher()
             publisher.init(conf=Scoped.get_scoped_conf(conf=self.conf,
                                                        scope=publisher.get_scope()))
@@ -89,7 +96,10 @@ class TestElasticsearchPublisher(unittest.TestCase):
                                 'KEY_DOESNOT_MATTER2': 'NO_VALUE2'})
         self.mock_es_client.indices.get_alias.return_value = {'test_old_index': 'DOES_NOT_MATTER'}
 
-        with patch('__builtin__.open', mock_open(read_data=mock_data)) as mock_file:
+        target = 'builtins.open'
+        if six.PY2:
+            target = '__builtin__.open'
+        with patch(target, mock_open(read_data=mock_data)) as mock_file:
             publisher = ElasticsearchPublisher()
             publisher.init(conf=Scoped.get_scoped_conf(conf=self.conf,
                                                        scope=publisher.get_scope()))

tests/unit/transformer/test_sql_to_test_usage.py

+import pytest
+import six
 import unittest
 from typing import no_type_check
 from mock import patch
 from pyhocon import ConfigFactory
 
 from databuilder.extractor.hive_table_metadata_extractor import HiveTableMetadataExtractor
-from databuilder.models.table_column_usage import TableColumnUsage, ColumnReader
 from databuilder.models.table_metadata import TableMetadata, ColumnMetadata
-from databuilder.transformer.sql_to_table_col_usage_transformer import SqlToTblColUsageTransformer
+if six.PY2:
+    from databuilder.models.table_column_usage import TableColumnUsage, ColumnReader
+    from databuilder.transformer.sql_to_table_col_usage_transformer import SqlToTblColUsageTransformer
 
 
+@pytest.mark.skipif(
+    six.PY3,
+    reason="Deactivated for Python3 because of issue [#40](https://github.com/lyft/amundsen/issues/40) (ANTLR for py3)"
+)
 class TestSqlToTblColUsageTransformer(unittest.TestCase):
 
     @no_type_check

tests/unit/usage/presto/test_column_usage_provider.py

 import logging
+import pytest
+import six
 import unittest
 
 from databuilder.sql_parser.usage.column import Column, Table, OrTable, remove_double_quotes
-from databuilder.sql_parser.usage.presto.column_usage_provider import ColumnUsageProvider
+if six.PY2:
+    from databuilder.sql_parser.usage.presto.column_usage_provider import ColumnUsageProvider
 
 
+@pytest.mark.skipif(
+    six.PY3,
+    reason="Deactivated for Python3 because of issue [#40](https://github.com/lyft/amundsen/issues/40) (ANTLR for py3)"
+)
 class TestColumnUsage(unittest.TestCase):
 
     def setUp(self):