Fixed the imports again

kayak · Jan 21, 2019 · 264fb0e · 264fb0e
1 parent 060d7b3
commit 264fb0e
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 16 deletions.
diff --git a/fireant/database/snowflake.py b/fireant/database/snowflake.py
@@ -62,15 +62,15 @@ def __init__(self, user='snowflake', password=None,
         self.warehouse = warehouse
 
     def connect(self):
-        from snowflake import connector as snowflake
-
-        return snowflake.connect(database=self.database,
-                                 account=self.account,
-                                 user=self.user,
-                                 password=self.password,
-                                 private_key=self._get_private_key(),
-                                 region=self.region,
-                                 warehouse=self.warehouse)
+        import snowflake
+
+        return snowflake.connector.connect(database=self.database,
+                                           account=self.account,
+                                           user=self.user,
+                                           password=self.password,
+                                           private_key=self._get_private_key(),
+                                           region=self.region,
+                                           warehouse=self.warehouse)
 
     def trunc_date(self, field, interval):
         trunc_date_interval = self.DATETIME_INTERVALS.get(str(interval), 'DD')

diff --git a/fireant/tests/database/test_snowflake.py b/fireant/tests/database/test_snowflake.py
@@ -24,9 +24,12 @@ def test_defaults(self):
         self.assertIsNone(snowflake.warehouse)
 
     def test_connect_with_password(self):
+        mock_snowflake = Mock(name='mock_snowflake')
+        mock_connector = mock_snowflake.connector
+
         # need to patch this here so it can be imported in the function scope
-        with patch('snowflake.connector') as mock_snowflake:
-            mock_snowflake.connect.return_value = 'OK'
+        with patch.dict('sys.modules', snowflake=mock_snowflake):
+            mock_connector.connect.return_value = 'OK'
 
             snowflake = SnowflakeDatabase(user='test_user',
                                           password='test_pass',
@@ -35,7 +38,7 @@ def test_connect_with_password(self):
             result = snowflake.connect()
 
         self.assertEqual('OK', result)
-        mock_snowflake.connect.assert_called_once_with(user='test_user',
+        mock_connector.connect.assert_called_once_with(user='test_user',
                                                        password='test_pass',
                                                        account='test_account',
                                                        database='test_database',
@@ -45,10 +48,13 @@ def test_connect_with_password(self):
 
     @patch('fireant.database.snowflake.serialization')
     def test_connect_with_pkey(self, mock_serialization):
+        mock_snowflake = Mock(name='mock_snowflake')
+        mock_connector = mock_snowflake.connector
+        mock_pkey = mock_serialization.load_pem_private_key.return_value = Mock(name='pkey')
+
         # need to patch this here so it can be imported in the function scope
-        with patch('snowflake.connector') as mock_snowflake:
-            mock_pkey = mock_serialization.load_pem_private_key.return_value = Mock(name='pkey')
-            mock_snowflake.connect.return_value = 'OK'
+        with patch.dict('sys.modules', snowflake=mock_snowflake):
+            mock_connector.connect.return_value = 'OK'
 
             snowflake = SnowflakeDatabase(user='test_user',
                                           private_key_data='abcdefg',
@@ -66,7 +72,7 @@ def test_connect_with_pkey(self, mock_serialization):
                                                                             backend=ANY)
 
         with self.subTest('connects with credentials'):
-            mock_snowflake.connect.assert_called_once_with(user='test_user',
+            mock_connector.connect.assert_called_once_with(user='test_user',
                                                            password=None,
                                                            account='test_account',
                                                            database='test_database',