Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

mocking snowflake connection

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:

  1. monkeypatch.setattr(connections.snowflake_connection,"connect",return_value = "")
  2. 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.

like image 379
Sheikh Usman Avatar asked Sep 05 '26 06:09

Sheikh Usman


1 Answers

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
like image 175
Hussain Bohra Avatar answered Sep 07 '26 18:09

Hussain Bohra



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!