Mocking subclasses created in __init__ - python

I have a database helper class that takes in a connection as an argument, and initializes another helper class in init
def generate_table_deprecation_lkp() -> None:
"""
run sql to generate lkp table that drives table deprecation workflow
"""
with RedshiftConnection() as conn:
helper = TableDeprecationHelper(conn)
helper.generate_table_deprecation_lkp()
class TableDeprecationHelper():
"""
Class that takes in RedshiftConnection and provides helper functions for table deprecation.
"""
def __init__(self, redshift_connection: RedshiftConnection) -> None:
self.redshift_connection = redshift_connection
self.table_manager = RedshiftTableManager(redshift_connection)
log.info("Initialized TableDeprecationHelper")
def generate_table_deprecation_lkp(self) -> None:
"""
generates transient lkp table used to drive deprecation workflow
"""
self.table_manager.truncate_table(*UNUSED_TABLES_LKP.split('.'))
self.redshift_connection.execute_db_command(FLAG_UNUSED_TABLES)
log.info("Generated table deprecation stage")
I'm trying to test a function that calls this helper using pytest, with code that looks like this
#pytest.fixture
def redshift_connection_patch():
with patch("RedshiftConnection") as redshift_connection_patch:
yield redshift_connection_patch
#pytest.fixture
def truncate_table_patch():
with patch.object(RedshiftTableManager, "truncate_table") as truncate_table_patch:
yield truncate_table_patch
def test_generate_table_deprecation_lkp(redshift_connection_patch, truncate_table_patch):
generate_table_deprecation_lkp()
truncate_table_patch.assert_called_with("admin", "table_deprecation_lkp")
redshift_connection_patch.execute_db_command.assert_called_with(FLAG_UNUSED_TABLES)
The truncate_table passes as expected. However, the redshift_connection_patch isn't mocking as expected, and I keep ending up with an error like this:
py::test_generate_table_deprecation_lkp Failed: [undefined]AssertionError: expected call not found.
Expected: execute_db_command("\ninsert into admin
I can see my db call on redshift_connection_patch.mock_calls, but it's not appearing in method calls, and the execute_db_command method is missing from the mock object altogether.

Related

How to test operations in a context manager using pytest

I have a database handler that utilizes SQLAlchemy ORM to communicate with a database. As part of SQLAlchemy's recommended practices, I interact with the session by using it as a context manager. How can I test what a function called inside the context manager using that context manager has done?
EDIT: I realized the file structure mattered due to the complexity in introduced. I re-structured the code below to more closely mirror what the end file structure will be like, and what a common production repo in my environment would look like, with code being defined in one file and tests in a completely separate file.
For example:
Code File (delete_things_from_table.py):
from db_handler import delete, SomeTable
def delete_stuff(handler):
stmt = delete(SomeTable)
with handler.Session.begin() as session:
session.execute(stmt)
session.commit()
Test File:
import pytest
import delete_things_from_table as dlt
from db_handler import Handler
def test_delete_stuff():
handler = db_handler()
dlt.delete_stuff(handler):
# Test that session.execute was called
# Test the value of 'stmt'
# Test that session.commit was called
I am not looking for a solution specific to SQLAlchemy; I am only utilizing this to highlight what I want to test within a context manager, and any strategies for testing context managers are welcome.
After sleeping on it, I came up with a solution. I'd love additional/less complex solutions if there are any available, but this works:
import pytest
import delete_things_from_table as dlt
from db_handler import Handler
class MockSession:
def __init__(self):
self.execute_params = []
self.commit_called = False
def execute(self, *args, **kwargs):
self.execute_params.append(["call", args, kwargs])
return self
def commit(self):
self.commit_called = True
return self
def begin(self):
return self
def __enter__(self):
return self
def __exit__(self, type, value, traceback):
pass
def test_delete_stuff(monkeypatch):
handler = db_handler()
# Parens in 'MockSession' below are Important, pass an instance not the class
monkeypatch.setattr(handler, Session, MockSession())
dlt.delete_stuff(handler):
# Test that session.execute was called
assert len(handler.Session.execute_params)
# Test the value of 'stmt'
assert str(handler.Session.execute_params[0][1][0]) == "DELETE FROM some_table"
# Test that session.commit was called
assert handler.Session.commit_called
Some key things to note:
I created a static mock instead of a MagicMock as it's easier to control the methods/data flow with a custom mock class
Since the SQLAlchemy session context manager requires a begin() to start the context, my mock class needed a begin. Returning self in begin allows us to test the values later.
context managers rely on on the magic methods __enter__ and __exit__ with the argument signatures you see above.
The mocked class contains mocked methods which alter instance variables allowing us to test later
This relies on monkeypatch (there are other ways I'm sure), but what's important to note is that when you pass your mock class you want to patch in an instance of the class and not the class itself. The parentheses make a world of difference.
I don't think it's an elegant solution, but it's working. I'll happily take any suggestions for improvement.

Mock an external service instantiated in python constructor

I'm working on unit tests for a service I made that uses confluent-kafka. The goal is to test successful function calls, exception errors, etc. The problem I'm running into is since I'm instantiating the client in the constructor of my service the tests are failing since I'm unsure how to patch a constructor. My question is how do I mock my service in order to properly test its functionality.
Example_Service.py:
from confluent_kafka.schema_registry import SchemaRegistryClient
class ExampleService:
def __init__(self, config):
self.service = SchemaRegistryClient(config)
def get_schema(self):
return self.service.get_schema()
Example_Service_tests.py
from unittest import mock
#mock.patch.object(SchemaRegistryClient, "get_schema")
def test_get_schema_success(mock_client):
schema_Id = ExampleService.get_schema()
mock_service.assert_called()
The problem is that you aren't creating an instance of ExampleService; __init__ never gets called.
You can avoid patching anything by allowing your class to accept a client maker as an argument (which can default to SchemaRegistryClient:
class ExampleService:
def __init__(self, config, *, client_factory=SchemaRegistryClient):
self.service = client_factory(config)
...
Then in your test, you can simply pass an appropriate stub as an argument:
def test_get_schema_success():
mock_client = Mock()
schema_Id = ExampleService(some_config, client_factory=mock_client)
mock_client.assert_called()
Two ways
mock entire class using #mock.patch(SchemaRegistryClient) OR
replace #mock.patch.object(SchemaRegistryClient, "get_schema") with
#mock.patch.object(SchemaRegistryClient, "__init__")
#mock.patch.object(SchemaRegistryClient, "get_schema")

Using __init__ method to call class main functionality

I needed to encapsulate functions related to single responsibility of parsing and emiting messages to API endpoint so I have created class Emitter.
class Emitter:
def __init__(self, message: str) -> None:
self.parsed = self.parse_message(message)
self.emit_message(self.parsed)
#staticmethod
def parse_message(msg: str) -> str:
... # parsing code
#staticmethod
def emit_message(msg: str) -> None:
... # emitting code
In order to emit a message I call a short-lived instance of that class with message passed as argument to __init__.
Emitter("my message to send")
__init__ itself directly runs all necessary methods to parse and emit message.
Is it correct to use __init__ to directly run the main responsibility of a class? Or should I use different solution like creating function that first instantiates the class and then calls all the necessary methods?
It looks like you're attempting to do the following at once:
initialize a class Emitter with an attribute message
parse the message
emit the message
IMO, a class entity is a good design choice since it allows granular control flow. Given that each step from above is discrete, I'd recommend executing them as such:
# Instantiate (will automatically parse)
e = Emitter('my message')
# Send the message
e.send_message()
You will need to redesign your class to the following:
class Emitter:
def __init__(self, message: str) -> None:
self.message = message
self.parsed = self.parse_message(message)
#staticmethod
def parse_message(msg: str) -> str:
... # parsing code
# This should be an instance method
def emit_message(msg: str) -> None:
... # emitting code
Also, your parse_message() method could be converted to a validation method, but that's another subject.

access python unittest magicmock return value

I'm using python 3.9.2 with unittest and mock to patch out a class.
My code under test instantiates an object of the class and mock returns a MagicMock object as the instance.
My question is, can I access that object from my test code?
I can see the call that instantiates the class in the mock_calls list, but cannot find a way of accessing the instance that is returned from that call.
The reason I need to access the instance is that my code under test attaches attributes to the instance rather than call methods on it. It is easy to test method calls, but is there a direct way to test attributes?
Upon investigation I found that there was only a single instance of a MagicMock being created and returned each time I instantiated my class. This behaviour was not convenient for me due to the attributes that I add to the class.
I created the following test aid to support my needs. This is not general-purpose but could be adapted for other circumstances.
class MockMyClass():
"""mock multiple MyClass instances
Note - the code under test must add a name attribute to each instance
"""
def __init__(self):
self.myclass = []
def factory(self, /, *args, **kwargs):
"""return a new instance each time called"""
new = mock.MagicMock()
# override __enter__ to enable the with... context manager behaviour
# for convenience in testing
new.__enter__ = lambda x: new
self.myclass.append(new)
return new
def __getitem__(self, key: str) -> None:
"""emulate a dict by returning the named instance
use as
mockmyclass['name'].assert_called_once()
or
with mockmyclass['name'] as inst:
inst.start().assert_called_once()
"""
# Important - the code under test gives the instance a name
# attribute and this relies on that attribute so is not
# general purpose
wanted = [t for t in self.myclass if t.name == key]
if not wanted:
names = [t.name for t in self.myclass]
raise ValueError(f'no timer {key} in {names}')
return wanted[0]
class TestBehaviour(unittest.TestCase):
def setUp(self):
self.mockmyclass = MockMyClass()
self.mocked = mock.patch(
'path-to-my-file.MyClass',
side_effect=self.mockmyclass.factory,
)
self.addCleanup(self.mocked.stop)
self.mocked = self.mocked.start()
def test_something(self):
# call code under test
# then test with
with self.mockmyclass['name-of-instance'] as inst:
inst.start.assert_called_once()
inst.stop.assert_called_once()
# or test with
self.mockmyclass['name-of-instance'].start.assert_called_once()

Mocking chained calls in Python

I am trying to test the following class using unittest and the mock library:
class Connection(object):
def __init__(self, cookie):
self.connect = None
self.session = Session()
self.session.load(cookie)
# do some stuff with self.session
self.some_info = self.session.data['the_info']
How could I test if when I create an instance of Connection, depending on the return of the Session instance, I assert if self.some_info is with the value I am expecting?
I wish to use the mock library. In its documentation I have an example of mocking chained calls (http://www.voidspace.org.uk/python/mock/examples.html#mocking-chained-calls), but it isn't very clear of how I can adapt it to my problem.
The Session.load(cookie) method sets some attributes in the Session instance. I would like to set this values fixed for my tests for every value of cookie.
Assume Connection is located in module package.module.connection
The following code should be how you would test your session:
import mock
class TestConnection(unittest.TestCase):
#mock.patch('package.module.connection.Session')
def test_some_info_on_session_is_set(self, fake_session):
fake_session.data = {'the_info': 'blahblah'}
cookie = Cookie()
connection = Connection(cookie)
self.assertEqual(connection.some_info, 'blahblah')
fake_session.load.assert_called_once_with(cookie)

Categories