Monkey Patching
At work I ran into a problem with ensuring that a Mongo test database was cleaned up between tests. As an ugly hack, we explicitly cleaned up after each fixture and needed to remember to do the same in tests that created new objects. This is burdensome and unmaintainable.
# an ugly solution to cleaning up after tests
@pytest.fixture()
def user():
id = ObjectID()
yield User.objects.create(id=id, name='jane doe')
User.objects.delete(id=id)
It doesn’t have to be this way. With pytest-django
, tests that access the database are marked with @pytest.mark.django_db
and pytest-django
ensures that the database is cleaned after each test. A mongo solution like this would solve my issue.
First attempt
My first solution was to create a fixture that would cleanup the database after each call.
connection = mongoengine.connect("testdb", host='mongomock://localhost')
@pytest.fixture()
def mongo():
"""Clear database after each test run"""
yield
connection.drop_database("testdb")
The problem with this solution is that each test that uses mongo must use this mongo
fixture. If you forgot to mark a test, you can leave data in your test database that could break other tests.
@pytest.mark.usefixtures("mongo")
def test_example():
User.objects.create(name='jane doe')
assert User.objects.count() == 1
A better way
The solution is to mimic the behavior of pytest-django
. In pytest-django
, if you attempt to access the database in a test that isn’t marked django_db
, an exception is raised. This mark enables pytest-django
to hook up a fixture that cleans up the database after the marked test runs.
How can we control access to our Mongo database to only tests that are marked explicitly? We need to patch our ORM.
Patching the ORM
I’m using mongoengine
for my ORM, so the location to mock is specific to this library, but the procedure should be similar for other libraries.
After some stepping through a call to User.objects.create
I found that each call depended on the _get_collection
classmethod in the parent Document class. If we block access to this method, we can disable database access.
Note that since we want to block a classmethod, we must access __dict__
to save our method.
Full example
from mongoengine.document import Document
def mock_get_collection(*args, **kwargs):
"""Mock method to disable database access"""
class MongoDBAccessNotAllowed(BaseException):
"""MongoDB access not allowed without fixture"""
raise MongoDBAccessNotAllowed()
# block access by default
# `_get_collection` is a classmethod so we must retrieve it from `__dict__` so
# we don't evaluate it via getattr.
original = Document.__dict__.get('_get_collection')
Document._get_collection = mock_get_collection
@pytest.fixture(autouse=True)
def mongo_access(request):
"""
Only allow mongo access to tests that request it.
This allows us to cleanup after tests that touch mongo.
"""
use_mongo = request.node.keywords.get("mongo")
if use_mongo:
# allow database access
Document._get_collection = original
yield
# disable access
Document._get_collection = mock_get_collection
# cleanup after ourselves
connection.drop_database(DB_NAME)
else:
# do nothing for tests that don't require access
yield
For our tests that we want to access the mongo database from we need to mark them with the decorator mongo
as we do in the following.
@pytest.mark.mongo
def test_example(user):
user.name = 'jane doe'
user.save()
assert user.id is not None
If we didn’t mark this test, MongoDBAccessNotAllowed
would be raised.