Mocking a Sqlalchemy session for pytest

Question

I don't know if this can be done but I'm trying to mock my db.session.save.

I'm using flask and flask-alchemy.

db.py

from flask_sqlalchemy import SQLAlchemy

db = SQLAlchemy()

The unit test

def test_post(self):
    with app.app_context():
        with app.test_client() as client:
            with mock.patch('models.db.session.save') as mock_save:
                with mock.patch('models.db.session.commit') as mock_commit:

                    data = self.gen_legend_data()
                    response = client.post('/legends', data=json.dumps([data]), headers=access_header)

                    assert response.status_code == 200
                    mock_save.assert_called()
                    mock_commit.assert_called_once()

And the method:

def post(cls):
    legends = schemas.Legends(many=True).load(request.get_json())

    for legend in legends:
        db.session.add(legend)

    db.session.commit()

    return {'message': 'legends saved'}, 200

I'm trying to mock the db.session.add and db.session.commit. I've tried db.session.save and legends.models.db.session.save and models.db.session.save. They all came back with the save error:

ModuleNotFoundError: No module named 'models.db.session'; 'models.db' is not a package

I don't get the error and I'm not sure how to solve it.

Or am I doing something that is totally wrong in wanting to mock a db.session?

Thanks. Desmond

Solution

The problem you're running into here is better served by restructuring your code so that it's more testable rather than mocking out every component, or otherwise making a (very) slow integration test. If you get in the habit of writing tests in that way, then over time you'll end up with a slow build that will take too long to run, and you'll end up with fragile tests (good talk on the subject of why fast tests are important here).

Let's take a look at this route:

def post(cls):
    legends = schemas.Legends(many=True).load(request.get_json())

    for legend in legends:
        db.session.add(legend)

    db.session.commit()

    return {'message': 'legends saved'}, 200

...and decompose it:

import typing
from flask import jsonify

class LegendsPostService:

    def __init__(self, json_args, _session=None) -> None:
        self.json_args = json_args
        self.session = _session or db.session

    def _get_legends(self) -> Legend:
        return schemas.Legends(many=True).load(self.json_args)

    def post(self) -> typing.List[typing.Dict[str, typing.Any]]:
        legends = self._get_legends()

        for legend in legends:
            self.session.add(legend)

        self.session.commit()
        return schemas.Legends(many=True).dump(legends)

def post(cls):
    service = LegendsPostService(json_args=request.get_json())
    service.post()
    return jsonify({'message': 'legends saved'})

Notice how we've isolated nearly all the points of failure from post into LegendsPostService, and further, we've removed all the flask internals from it as well (no global request objects floating around, etc). We've even given it the ability to mock out session if we need to for testing later on.

I would recommend you focus your testing efforts on writing test cases for LegendsPostService. Once you've got excellent tests for LegendsPostService, decide if you believe that even more test coverage will add value. If you do, then consider writing one simple integration test for post() to tie it all together.

The next thing you need to consider is how you want to think about SQLAlchemy objects in tests. I recommend just using FactoryBoy for auto-creating "mock" models for you. Here's a full application example for how to setup flask / sqlalchemy / factory-boy in this way: How do I produce nested JSON from database query with joins? Using Python / SQLAlchemy

Here's how I'd write a test for LegendsPostService (apologies as this is a bit hasty and doesn't perfectly represent what you're trying to do - but you should be able to adjust these tests for your use case):

from factory.alchemy import SQLAlchemyModelFactory

class ModelFactory(SQLAlchemyModelFactory):
    class Meta:
        abstract = True
        sqlalchemy_session = db.session

# setup your factory for Legends:
class LegendsFactory(ModelFactory):
    logo_url = factory.Faker('image_url')
    class Meta(ModelFactory.Meta):
        model = Legends


from unittest.mock import MagicMock, patch


# neither of these tests even need a database connection!
# so you should be able to write HUNDREDS of similar tests
# and you should be able to run hundreds of them in seconds (not minutes)

def test_LegendsPostService_can_init():
    session = MagicMock()
    service = LegendsPostService(json_args={'foo': 'bar'}, _session=session)
    assert service.session is session
    assert service.json_args['foo'] == 'bar'


def test_LegendsPostService_can_post():
    session = MagicMock()
    service = LegendsPostService(json_args={'foo': 'bar'}, _session=session)

    # let's make some fake Legends for our service!
    legends = LegendsFactory.build_batch(2)

    with patch.object(service, '_get_legends') as _get_legends:
        _get_legends.return_value = legends
        legends_post_json = service.post()

    # look, Ma! No database connection!
    assert legends_post_json[0]['image_url'] == legends[0].image_url

I hope that helps!

Comments

[NotMe]

Thanks Stephen. That helped a lot. Just a question, the line session = MagicMock() is mocking the db.session, correctly. So when the codes get to the self.session.add(), it will just use the MagicMock() object and not call the actual session.

[Stephen Fuhry]

Yes, that's accurate. When you mock session, you are writing a unit test and not a system / integrated test. If you feel the need to have an integrated test to make sure all the sqlalchemy / database behavioral interactions are working as expected, you can write an integration test for that.

[Stephen Fuhry]

I recommend writing as few integrated tests as possible (i.e. ones that don't mock sqlalchemy & talk to the database). Why? Several reasons, but one of the more subtle reasons is one that you might not notice until a couple years into your project.. Consider that running 1000 unit tests should take less than 30 seconds. running 1000 integration tests, however, may take 30mins+