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.