How to mock mongodb when it is called from another function? - python

I need help while mocking mongodb. I am using mongomock to mock mongodb.
My project structure is:
-- my_mongo.py
-- code.py
-- my_test.py
my_mongo.py has :
from pymongo import MongoClient
def get_db():
client = MongoClient(os.environ['MONGODB_URI'])
db = client['my_db']
return db
def insert(id, data):
return get_db().results.insert_one(
{
"id": id,
"data":df.to_json(),
}).id
and code.py has
import my_mongo
def action():
#do somethings
my_mongo.insert(id, data)
and my_test.py has
import mongomock
import my_mongo
from unittest import mock
with patch.object(my_mongo.get_db().client, "client", mongomock.MongoClient()):
import code
def test_action_1():
my_mongo.insert = mock.Mock(return_value=1)
code.action()
def test_action_2():
with patch.object(my_mongo.get_db(), "get_db", mongomock.MongoClient().db):
code.action()
It throws pymongo.errors.ServerSelectionTimeoutError for both tests. So, It still goes into the insert_one() method in my_mongo.py.
I expect in test_action_1 my_mongo.insert returns 1, but it doesn't.
What am I missing?

I'm not entirely sure what mongomock is for, but it looks like it's for mocking an entire mongo database and not actually using python mocking. I'm going to answer without including mongomock since I don't think you really need it, so you can take that for what it's worth.
There were a few issues:
Calling patch.object will patch the given method on whatever object you give it. If you call get_db in the test, then code.action calls get_db, those are 2 different objects. Maybe this works? But I'm skeptical, so I just changed it.
Don't use code as your module name. That's already a module included with python.
code.action was missing args and a return statement.
You'll also notice that I changed how and what was being mocked to illustrate different ways to accomplish the mocking. Test 1 mocks the insert call with a function decorator. Test 2 mocks the get_db call with a contextmanager. Either is correct, just showing that you have options.
Here's the finished product:
my_mongo.py:
from pymongo import MongoClient
def get_db():
client = MongoClient(os.environ['MONGODB_URI'])
db = client['my_db']
return db
def insert(id, data):
return get_db().results.insert_one({"id": id, "data":data.to_json()}).id # df was undefined, updated to data
my_code.py:
import my_mongo
# I added id and data args. They were undefined
def action(id, data):
return my_mongo.insert(id, data) # I added a return here
my_test.py
from unittest import mock
import my_code
# I removed the contextmanager import. Nothing is being evaluated here that would
# need to be patched, so I'm pretty certain it has no effect
#mock.patch('my_mongo.insert')
def test_action_1(mock_insert):
expected_id = 1
mock_insert.return_value = expected_id
ret = my_code.action(expected_id, mock.Mock())
assert ret == expected_id
def test_action_2():
with mock.patch('my_mongo.get_db') as mock_get_db:
expected_id = 'some id'
mock_db = mock.Mock()
mock_db.results.insert_one.return_value.id = expected_id
mock_get_db.return_value = mock_db
ret = my_code.action(expected_id, mock.Mock())
assert ret == expected_id

That line of code for patching mongodb is wrong. Instead of using patch.object(my_mongo.get_db(), "get_db", mongomock.MongoClient().db), you should use patch.object("my_mongo.get_db", return_value=mongomock.MongoClient()['my_db']).
Following is the complete runable code for your example:
my_test.py
import mongomock
from unittest.mock import patch
import my_code
import my_mongo
def test_action_2():
mocked_mongo = mongomock.MongoClient()
with patch("my_mongo.get_db", return_value=mongomock.MongoClient()['my_db']):
my_code.action()
assert mocked_mongo.my_db.results.count_documents({'id': 'some_id'}) == 1
my_mongo.py
from pymongo import MongoClient
def get_db():
client = MongoClient(os.environ['MONGODB_URI'])
db = client['my_db']
return db
def insert(id, data):
return get_db().results.insert_one(
{
"id": id,
"data": data,
})
my_code.py
import my_mongo
def action():
#do somethings
return my_mongo.insert('some_id', '{"a": 3}')

Related

Python mock.patch() not mocking imported class module

I have a some code that looks like this (abridged):
## ./module_1/lambda_function.py
from shared.graphql_helpers import GraphQLClient
gql_client = GraphQLClient('http://host_url.test/')
post_event_to_consumer(event):
response = gql_client.make_query(mutation, { 'data': event, {})
lambda_hanbdler(event, context):
post_event_to_consumer(event['detail']
The module imported is such:
## ./shared/graphql_helpers.py
import requests
class GraphQLClient:
def __init__(self, url):
self.host_url = url
def make_query(self, query, variables, headers):
request = requests.post(self.host_url, json={'query': query, 'variables': variables}, headers=headers)
if request.status_code != 200:
raise Exception()
elif request.json()['errors']:
raise Exception()
else:
return request.json()['data']
The project structure is such:
module_1
__init__.py
lambda_function.py
lambda_function_test.py
shared
__init__.py
graphql_helpers.py
My problem is that I am trying to patch GraphQLClient in lambda_function_test.py but it does not seem to be working.
My unit test looks like this:
## ./module_1/lambda_function_test.py
import os
from unittest import mock, TestCase
class TestLambdaFunction(TestCase):
#mock.patch('module_1.lambda_function.GraphQLClient')
def test_post_event_to_compplex(self, mock_gql_client):
mock_gql_client = mock.Mock()
mock_gql_client.make_query.return_value = {}
result = post_event_to_consumer(json_data['detail'])
assert result == None
From what I read online, patch mocks the import in the SUT's namespace, however, when I run the test, it imports the original GraphQLClient and throws a ConnectionError. I also tried removing module_1 from the path but that didn't work either.
Not sure what's incorrect here, but any help would be appreciated.

How do you populate package variables when testing a class's method in Pytest?

I'm writing a test suite for a class that needs variables from a util package I have defined:
# util/__init__.py
codes = None
def load_codes():
"""Populates util.codes with the contents of codes.yml"""
with open('codes.yml', 'r') as f:
global codes
codes = yaml.safe_load(f)
The class under test, Filter, is a plugin for a larger app; normally util.load_codes() is called on startup of the app, so when Filter needs to read util.codes, it is guaranteed to be populated. Filter should not be concerned with loading the file into the package var.
For the test itself, I'm loading the class instance as a fixture:
# tests/filter.py
import pytest
from unittest.mock import AsyncMock
from plugins.filter import Filter
#pytest.fixture
def filterplugin():
bot = AsyncMock()
return Filter(bot=bot)
#pytest.mark.asyncio
class TestFilterLogic:
async def test_basic(self, filterplugin):
msg = AsyncMock()
output = await filterplugin.filter_handler(msg)
assert output == True
Filter itself is large and complicated, but a MVP that demonstrates the problem is:
# plugins/filter.py
from util import codes
class Filter():
async def filter_handler(self):
print(codes.keys())
return True
This test fails with an AttributeError when filter_handler calls .keys() on util.codes; it's None.
I attempted to fix this by calling importing util and calling util.load_codes() both in the fixture and in the test itself, to no effect. The test still fails and indicates that the codes var in util is None.
How do I correctly populate this package variable so my class can read it when being unit tested independently of the rest of the larger app?
when Filter needs to read util.codes
The implementation shown does not read util.codes. It reads a local reference to the value of util.codes at import time, i.e. a local reference to None. That local variable will not be modified by load_codes().
There are a few ways to resolve this.
Initialize codes as a dictionary and modify that instance (not reassign) in load_codes().
# util/__init__.py
# codes = None # -
codes = {} # +
def load_codes():
"""Populates util.codes with the contents of codes.yml"""
with open('codes.yml', 'r') as f:
global codes
# codes = yaml.safe_load(f) # -
codes.clear() # +
codes.update(yaml.safe_load(f)) # +
Import the module and reference the module variable as an attribute of the module.
# plugins/filter.py
# from util import codes # -
import util # +
class Filter():
async def filter_handler(self):
# print(codes.keys()) # -
print(util.codes.keys()) # +
return True
Import the module variable as a local variable lazily.
# plugins/filter.py
# from util import codes # -
class Filter():
async def filter_handler(self):
from util import codes # +
print(codes.keys())
return True
Import Filter lazily, assuming load_codes() is called before the fixture.
# tests/filter.py
import pytest
from unittest.mock import AsyncMock
# from plugins.filter import Filter # -
#pytest.fixture
def filterplugin():
from plugins.filter import Filter # +
bot = AsyncMock()
return Filter(bot=bot)

Unit test for Python AWS Lambda: mock function before module is imported

I'm trying to write unit tests for my aws lambda function written in python 3.9. I tried different things to mock the get_object function that makes calls to the S3. I wanted to focus only on the calculate method, to verify if I'm getting correct results of an calculation.
When I try to run the following approach I'm getting credential errors about boto3
python -m unittest tests/app-test.py
...
botocore.exceptions.NoCredentialsError: Unable to locate credentials
Is there a way to import the calculate method from app.py and mock call to the get_object fn?
directory:
functions:
- __init__.py
- app.py
tests:
- __init__.py
- app-test.py
lambda function app.py:
import json
import boto3
def get_object():
s3 = boto3.client('s3')
response = s3.get_object(Bucket='mybucket', Key='object.json')
content = response['Body'].read().decode('utf-8')
return json.loads(content)
stops = get_object()
def lambda_handler(event, context):
params = event['queryStringParameters']
a = int(params['a'])
b = int(params['b'])
result = calculate(a, b)
return {
'statusCode': 200,
'body': json.dumps(result)
}
def calculate(a, b):
return a + b
unit test app-test.py:
import unittest
from unittest import mock
with mock.patch('functions.app.get_object', return_value={}):
from functions.app import calculate
class TestCalculation(unittest.TestCase):
def test_should_return_correct_calculation(self):
# when
result = calculate(1, 2)
# then
self.assertEqual(3, result)
I was able to fix the issue. The biggest obstacle was to mock the boto3 in the app.py. I did this, by mocking the whole boto3 module before it's imported. Here's the code of app-test.py
import sys
from io import BytesIO
from json import dumps
from unittest import TestCase, main
from unittest.mock import Mock
from botocore.stub import Stubber
from botocore.session import get_session
from botocore.response import StreamingBody
# prepare mocks for boto3
stubbed_client = get_session().create_client('s3')
stubber = Stubber(stubbed_client)
# mock response from S3
body_encoded = dumps({'name': 'hello world'}).encode()
body = StreamingBody(BytesIO(body_encoded), len(body_encoded))
stubbed.add_response('get_object', {'Body': body})
stubber.activate()
# add mocks to the real module
sys.modules['boto3'] = Mock()
sys.modules['boto3'].client = Mock(return_value=stubbed_client)
# Import the module that will be tested
# boto3 should be mocked in the app.py
from functions.app import calculate
class TestCalculation(TestCase):
def test_should_return_correct_calculation(self):
# when
result = calculate(1, 2)
# then
self.assertEqual(3, result)

Python mocking not working for shutil.rmtree

I have a class with a method which uses shutil.rmtree to remove some files if a param is passed as true, How to mock this behavior so as other tests don't break which needs these files.
My class looks like this -
class FileConverter(object):
def __init__(self, path_to_files):
self._path_to_files = path_to_files
def convert_files(self, rmv_src=False):
doStuff()
if rmv_src:
shutil.rmtree(self.__path_to_files)
def doStuff():
# does some stuff
Now my tests look like -
class TestFileConverter(unittest.TestCase):
def test_convert_success(self):
input_dir = 'resources/files'
file_converter = FileConverter(input_dir)
file_converter.convert_files()
# assert the things from doStuff
#mock.patch('shutil.rmtree')
def test_convert_with_rmv(self, rm_mock):
input_dir = 'resources/files'
file_converter = FileConverter(input_dir)
file_converter.convert_files(True)
self.assertEquals(rm_mock, [call(input_dir)])
Now when I run this testsuite the test with rmv gives me assertionError
<MagicMock name='rmtree' id='139676276822984'> != [call('resources/images')]
and the first test gives me file not found error since the mock did not work and the rmv source test removed the file
FileNotFoundError: [Errno 2] No such file or directory: 'resources/images'
If I comment out the second test with rmv_source true then my first test works fine.
What am I doing wrong here?
Your module has already imported shutil.rmtree so mocking it later in the test suite won't do anything.
You need to mock the module when you import FileConverter, not afterwards.
import sys
from mock import MagicMock
sys.modules['shutil'] = MagicMock()
# and/or
sys.modules['shutil.rmtree'] = MagicMock()
import FileConverter
If you still need to use shutil in your test code, then first import it using an alias,and use that when you need the 'real' module:
import sys
from mock import MagicMock
import shutil as shutil_orig
sys.modules['shutil'] = MagicMock()
import shutil
print(type(shutil_orig.rmtree))
# <class 'function'>
print(type(shutil.rmtree))
# <class 'mock.mock.MagicMock'>
The original post should work except the call(input_dir) did not work for me
#mock.patch('shutil.rmtree')
def test_convert_with_rmv(self, rm_mock):
input_dir = 'resources/files'
rm_mock.return_value = 'REMOVED'
file_converter = FileConverter(input_dir)
file_converter.convert_files(True)
rm_mock.assert_called_with(input_dir)
self.assertEqual(rm_mock.return_value, 'REMOVED')
The test_convert_with_rmv has no way removed the input_dir, it probably never created in the first place. You could assert this statement in each test before and after convert_files called:
self.asserTrue(os.path.isdir(input_dir))

Mocking two functions with patch for a unit test

I have a function I want to unit test contains calls two other functions. I am unsure how can I mock both functions at the same time properly using patch. I have provided an example of what I mean below. When I run nosetests, the tests pass but I feel that there must be a cleaner way to do this and I do not really Understand the piece regarding f.close()...
The directory structure looks like this:
program/
program/
data.py
tests/
data_test.py
data.py:
import cPickle
def write_out(file_path, data):
f = open(file_path, 'wb')
cPickle.dump(data, f)
f.close()
data_test.py:
from mock import MagicMock, patch
def test_write_out():
path = '~/collection'
mock_open = MagicMock()
mock_pickle = MagicMock()
f_mock = MagicMock()
with patch('__builtin__.open', mock_open):
f = mock_open.return_value
f.method.return_value = path
with patch('cPickle.dump', mock_pickle):
write_out(path, 'data')
mock_open.assert_called_once_with('~/collection', 'wb')
f.close.assert_any_call()
mock_pickle.assert_called_once_with('data', f)
Results:
$ nosetests
.
----------------------------------------------------------------------
Ran 1 test in 0.008s
OK
You can simplify your test by using the patch decorator and nesting them like so (they are MagicMock objects by default):
from unittest.mock import patch
#patch('cPickle.dump')
#patch('__builtin__.open')
def test_write_out(mock_open, mock_pickle):
path = '~/collection'
f = mock_open.return_value
f.method.return_value = path
write_out(path, 'data')
mock_open.assert_called_once_with('~/collection', 'wb')
mock_pickle.assert_called_once_with('data', f)
f.close.assert_any_call()
Calls to a MagicMock instance return a new MagicMock instance, so you can check that the returned value was called just like any other mocked object. In this case f is a MagicMock named 'open()' (try printing f).
In addition to the response #Matti John you can also use patch inside function test_write_out:
from mock import MagicMock, patch
def test_write_out():
path = '~/collection'
with patch('__builtin__.open') as mock_open, \
patch('cPickle.dump') as mock_pickle:
f = mock_open.return_value
...
As of Python 3.10 you can do use Parenthesized Context Managers like this
from unittest.mock import patch
def test_write_out():
with (
patch('cPickle.dump'),
patch('__builtin__.open') as open_mock, # example of using `as`
):
yield
Here's a simple example on how to test raising ConflictError in create_collection function using mock:
import os
from unittest import TestCase
from mock import patch
from ..program.data import ConflictError, create_collection
class TestCreateCollection(TestCase):
def test_path_exists(self):
with patch.object(os.path, 'exists') as mock_method:
mock_method.return_value = True
self.assertRaises(ConflictError, create_collection, 'test')
Please, also see mock docs and Michael Foord's awesome introduction to mock.

Categories

Resources