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-django ensures that the database is cleaned after each test. A mongo solution like this would solve my issue.
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, 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
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.
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.