模拟雪花连接

Question

我在 python 中有一个 SnowflakeApi class，它只是作为 SnowflakeConnection class 之上的包装器。我的 SnowflakeApi 是

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()

SnowflakeApi 的命名空间是 connection.snowflake_connection.SnowflakeApi（即我在名为 connections 的文件夹中有 snowflake_connection.py）

我想使用 pytest 和 unittest.mock 为此 class 编写单元测试。问题是我想模拟 'connect' 以便返回 MagicMock 对象并且不进行数据库调用。到目前为止我已经尝试过：

monkeypatch.setattr(connections.snowflake_connection,"connect",return_value = "")
将我原来的 class 更改为仅导入雪花。然后我创建了一个模拟对象并使用了 monkeypatch.setattr(snowflake_connection,"snowflake",my_mock_snowflake)。那也没用

简而言之，我尝试了其他一些方法，但没有任何效果。我只想模拟雪花连接，因此不会进行实际的数据库调用。

Answer 1

使用存根和依赖注入

from ... import SnowflakeApi

def some_func(*args, api=None, **kwargs):
    api = api or SnowflakeApi(...)
    conn = api.get_connection()
    # Do some work
    return result

你的测试

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 2

使用 unittest.mock 和修补连接的示例：

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')

如果您正在测试使用您的 SnowFlakeApi 包装器的其他类，那么您应该使用相同的方法，但在那些测试中修补 SnowFlakeApi 本身。

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()

另请注意，如果您使用的是 Python 2，则需要 pip install mock 然后 from mock import patch。

Answer 3

这是我们使用 python mock and patch 模拟 snowflake connector, cursor and fetch_all 的另一种方式。

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 4

以下解决方案对我有用。

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()

#模拟连接

def mocked_sf_connect(self, **kwargs):
    print("Connection Successfully Established")
    return True

Answer 5

使用 cursor、execute 和 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

在单元测试中，您模拟了预期的方法调用 - cursor()、execute()、 fetchone() 并定义 return 值以跟进定义的模拟链。

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()

模拟雪花连接

mocking snowflake connection

pytest

python-unittest

python-unittest.mock

snowflake-cloud-data-platform