mocking snowflake connection - pytest

Question

I have a SnowflakeApi class in python which just works as a wrapper on top of the SnowflakeConnection class. My SnowflakeApi is
import logging
import os
from snowflake.connector import connect
class SnowflakeApi(object):
"""
Wrapper to handle snowflake connection
"""
def __init__(self, account, warehouse, database, user, pwd):
"""
Handles snowflake connection. Connection must be closed once it is no longer needed
:param account:
:param warehouse:
:param database:
"""
self.__acct = self._account_url(account)
self.__wh = warehouse
self.__db = database
self.__connection = None
self.__user = user
self.__pwd = pwd
def __create_connection(self):
try:
# set the proxy here
conn = connect(
account=self.__acct
, user=self.__user
, password=self.__pwd
, warehouse=self.__wh
, database=self.__db
)
return conn
except:
raise Exception(
"Unable to connect to snowflake for user: '{0}', warehouse: '{1}', database: '{2}'".format(
self.__user, self.__wh, self.__db))
def get_connection(self):
"""
Gets a snowflake connection. If the connection has already been initialised it is returned
otherwise a new connection is created
:param credentials_func: method to get database credentials.
:return:
"""
try:
if self.__connection is None:
self.__connection = self.__create_connection()
return self.__connection
except:
raise Exception("Unable to initalise Snowflake connection")
def close_connection(self):
"""
Closes snowflake connection.
:return:
"""
self.__connection.close()
Namespace for SnowflakeApi is connection.snowflake_connection.SnowflakeApi (i.e. i have snowflake_connection.py in a folder called connections)
I want to write unit tests for this class using pytest and unittest.mock. The problem is I want to mock 'connect' so that a MagicMock object is returned and no database call is made. So far I have tried:
monkeypatch.setattr(connections.snowflake_connection,"connect",return_value = "")
Changed my original class to just import snowflake. I then created a mock object and used monkeypatch.setattr(snowflake_connection,"snowflake",my_mock_snowflake). That didn't work either
In short, I have tried a couple of other things but nothing has worked. All I want to do is mock snowflake connection so no actual database call is made.

Answer 1

Here is another way where we are mocking snowflake connector, cursor and fetch_all using python mock and patch.
import mock
import unittest
from datetime import datetime, timedelta
import feed_daily_report
class TestFeedDailyReport(unittest.TestCase):
#mock.patch('snowflake.connector.connect')
def test_compare_partner(self, mock_snowflake_connector):
tod = datetime.now()
delta = timedelta(days=8)
date_8_days_ago = tod - delta
query_result = [('partner_1', date_8_days_ago)]
mock_con = mock_snowflake_connector.return_value
mock_cur = mock_con.cursor.return_value
mock_cur.fetchall.return_value = query_result
result = feed_daily_report.main()
assert result == True

Answer 2

An example using unittest.mock and patching the connection:
from unittest import TestCase
from unittest.mock import patch
from connection.snowflake_connection import SnowflakeApi
class TestSnowFlakeApi(TestCase):
#patch('connection.snowflake_connection.connect')
def test_get_connection(self, mock_connect)
api = SnowflakeApi('the_account',
'the_warehouse',
'the_database',
'the_user',
'the_pwd')
api.get_connection()
mock_connect.assert_called_once_with(account='account_url', # Will be the output of self._account_url()
user='the_user',
password='the_pwd',
warehouse='the_warehouse',
database='the_database')
If you're testing other classes that use your SnowFlakeApi wrapper, then you should use the same approach, but patch the SnowFlakeApi itself in those tests.
from package.module.SomeClassThatUsesSnowFlakeApi
class TestSomeClassThatUsesSnowFlakeApi(TestCase):
#patch('package.module.SnowFlakeApi')
def test_some_func(self, mock_api):
instance = SomeClassThatUsesSnowFlakeApi()
instance.do_something()
mock_api.assert_called_once_with(...)
mock_api.return_value.get_connection.assert_called_once_with()
Also note that if you're using Python 2, you will need to pip install mock and then from mock import patch.

Answer 3

Using stubbing and dependency injection
from ... import SnowflakeApi
def some_func(*args, api=None, **kwargs):
api = api or SnowflakeApi(...)
conn = api.get_connection()
# Do some work
return result
Your test
class SnowflakeApiStub(SnowflakeApi)
def __init__(self):
# bypass super constructor
self.__connection = MagicMock()
def test_some_func():
stub = SnowflakeApiStub()
mock_connection = stub.__connection
mock_cursor = mock_connection.cursor.return_value
expect = ...
actual = some_func(api=stub)
assert expect == actual
assert mock_cursor.execute.called

Answer 4

An example using cursor, execute, and fetchone.
import snowflake.connector
class AlongSamePolly:
def __init__(self, conn):
self.conn = conn
def row_count(self):
cur = self.conn.cursor()
query = cur.execute('select count(*) from schema.table;')
return query.fetchone()[0] # returns (12345,)
# I like to dependency inject the snowflake connection object in my classes.
# This lets me use Snowflake Python Connector's built in context manager to
# rollback any errors and automatically close connections. Then you don't have
# try/except/finally blocks everywhere in your code.
#
if __name__ == '__main__':
with snowflake.connector.connect(user='user', password='password') as con:
same = AlongSamePolly(con)
print(same.row_count())
# => 12345
In the unittests you mock out the expected method calls - cursor(), execute(),
fetchone() and define the return value to follow up the chain of defined mocks.
import unittest
from unittest import mock
from along_same_polly import AlongSamePolly
class TestAlongSamePolly(unittest.TestCase):
def test_row_count(self):
with mock.patch('snowflake.connector.connect') as mock_snowflake_conn:
mock_query = mock.Mock()
mock_query.fetchone.return_value = (123,)
mock_cur = mock.Mock()
mock_cur.execute.return_value = mock_query
mock_snowflake_conn.cursor.return_value = mock_cur
same = AlongSamePolly(mock_snowflake_conn)
self.assertEqual(same.row_count(), 123)
if __name__ == '__main__':
unittest.main()

Answer 5

The following Solution Worked for me.
def test_connect(env_var_setup, monkeypatch):
monkeypatch.setattr(snowflake.connector.connection.SnowflakeConnection,
"connect", mocked_sf_connect
)
# calling snowflake connector method
file_job_map(env_var_setup).connect()
#mocked connection
def mocked_sf_connect(self, **kwargs):
print("Connection Successfully Established")
return True