Testing Flask SQLAlchemy database with pytest

Making tables accessible with create_all()

So far we haven’t put any code related to the database in our app fixture. The tutorial has done so in BaseTestCase. If we rewrite the test_add_user for pytest and try to run it, we get an error.

# project/tests/test_users.py
def test_add_user(app):
    """Ensure a new user can be added to the database."""
    with app.test_client() as client:
        response = client.post(
            '/users',
            data=json.dumps(dict(
                username='michael',
                email='michael@realpython.com'
            )),
            content_type='application/json',
        )
        data = json.loads(response.data.decode())
        assert response.status_code == 201
        assert 'michael@realpython.com was added!' in data['message']
        assert 'success' in data['status']

$ docker-compose run users-service python manage.py test_app
[...]
======= FAILURES ======
____ test_add_user ____
[...]
sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) relation "users" does not exist

So our app fixture knows about sql alchemy but hasn’t created the tables needed for our user model. From the Flask-SQLAlchemy documentation:

What it (init_db) does is prepare the application to work with SQLAlchemy. However that does not now bind the SQLAlchemy object to your application.

Let’s jump in an python shell and try to see how we can bind our db object to our app.

$ docker exec -ti users-service bash
root@910656bc5f75:/usr/src/app# python3
>>> from project import create_app, db
>>> app = create_app()
>>> app
<Flask 'project'>
>>> db
<SQLAlchemy engine=None>

So create_app() gives us an app object, but the imported db isn’t connected to it out-of-the-box (it has no engine). If we jump into the app context, the engine gets populated.

>>> app_ctx = app.app_context()
>>> app_ctx.push()  # jump into the app context
>>> db
<SQLAlchemy engine='postgres://postgres:postgres@users-db:5432/users_dev'>

However, db still doesn’t have tables and running tests at this stage would keep saying that relation users does not exist. That’s why we need create_table().

>>> db.engine.table_names()  # Check the tables currently on the engine
[]                           # no table found
>>> db.create_all()          # Create the tables according to defined models
>>> db.engine.table_names()
['users']                    # Now table 'users' is found

We can now update our app fixture:

# project/tests/conftest.py
import pytest
from project import create_app, db
@pytest.fixture
def app():
    app = create_app()
    app.config.from_object('project.config.TestingConfig')
    with app.app_context():   
        # alternative pattern to app.app_context().push()
        # all commands indented under 'with' are run in the app context 
        db.create_all()
        return app

Cleaning database with drop_all()

Let’s check the impact of our tests on the database. Before running tests:

$ docker exec -ti $(docker ps -aqf "name=users-db") psql -U postgres

psql (9.6.3)
postgres=# \c users_dev
You are now connected to database "users_dev" as user "postgres".
users_dev=# \dt
No relations found.
users_dev=# \c users_test
You are now connected to database "users_test" as user "postgres".
users_test=# \dt
No relations found.

After running tests:

# Running pytest
$ docker-compose run users-service python manage.py test_app
[... Output of pytest ...]
========= 5 passed in 0.17 seconds =========

# Checking the database
$ docker exec -ti $(docker ps -aqf "name=users-db") psql -U postgres

psql (9.6.3)
postgres=# \c users_dev
You are now connected to database "users_dev" as user "postgres".
users_dev=# \dt
No relations found.
users_dev=# \c users_test
You are now connected to database "users_test" as user "postgres".
users_test=# \dt
users_test=# \dt
         List of relations
 Schema | Name  | Type  |  Owner   
--------+-------+-------+----------
 public | users | table | postgres
(1 row)

users_test=# SELECT * FROM users;
 id | username |         email          | active |         created_at         
----+----------+------------------------+--------+----------------------------
  1 | michael  | michael@realpython.com | f      | 2017-07-04 09:23:43.34457

We can see that the users table was successfully created in the users_test database, which we selected in conftest.py. A user was also successfully inserted. However, if we run the tests again and recheck the table, we can see how trouble is starting to creep in:

# Running test 
[...]

# Checking the database
[...]
users_test=# SELECT * FROM users;
 id | username |         email          | active |         created_at         
----+----------+------------------------+--------+----------------------------
  1 | michael  | michael@realpython.com | f      | 2017-07-04 09:23:43.34457
  2 | michael  | michael@realpython.com | f      | 2017-07-04 09:24:55.972571
(2 rows)

We are not cleaning up the database after our tests. To do so, we need to add drop_all() as a tear-down action for our app fixture (after yield). This will not only empty the table rows, but also delete the table itself:

# project/tests/conftest.py
import pytest
from project import create_app, db
@pytest.fixture
def app():
    app = create_app()
    app.config.from_object('project.config.TestingConfig')
    with app.app_context():   
        db.create_all()
        yield app   # Note that we changed return for yield, see below for why
        db.drop_all()

For the first time, we added a command (db.drop_all()) that needed to be executed after the test using the fixture. Previously we used return to get the app out of the fixture. But using return means ending the function. That’s where yield comes to the rescue. Unlike unittest, pytest does not put setup and teardown code in dedicated function. Everything that comes before return/yield is setup code, everything that comes after yield is teardown code. Documentation is here:

By using a yield statement instead of return, all the code after the yield statement serves as the teardown code.

If you try to run tests again, you will see that the database is left clean.

Avoid locking postgres with db.session.remove()

This is the part I still have trouble understanding. Using the fixture above, pytest started hanging indefinitely at random test (usually at tests that touched the database several times, but not always). When it happened, I could not even stop pytest and had to restart the container.

$ docker-compose run users-service python manage.py test_app
===== test session starts =====
[...]
project/tests/test_configs.py ...
project/tests/test_users.py ..

hanging... hanging... hanging...

From this SO question, I got the confirmation that postgres might be locked. Using the commands listed on devopsderek.com, it is possible to look at the sessions active while pytest is locked. Below we can see that one is blocked with the state “idle in transaction”.

# Checking the database while pytest is hanging
$ docker exec -ti $(docker ps -aqf "name=users-db") psql -U postgres
postgres=# SELECT * FROM pg_stat_activity;

datid |  datname   | pid | ... | state               | ... | query
------+------------+-----+-----+---------------------+-----+-
[... lots of sessions ...]
16386 | users_test | 200 | ... | idle in transaction | ... | SELECT users.id...

Killing the session manually in postgres did also kill pytest without needing to restart the container.

# Still in psql
postgres=# SELECT pg_terminate_backend(200);
 pg_terminate_backend 
----------------------
 t
(1 row)

# In the terminal running pytest
psycopg2.OperationalError: terminating connection due to administrator command
server closed the connection unexpectedly
    This probably means the server terminated abnormally
    before or while processing the request.
$ # free to type again

Still following SO advice, you can explicitly close the connection after each test, which solves the issue.

# project/tests/conftest.py
import pytest
from project import create_app, db
@pytest.fixture
def app():
    app = create_app()
    app.config.from_object('project.config.TestingConfig')
    with app.app_context():   
        db.create_all()
        yield app  
        db.session.remove()  # looks like db.session.close() would work as well
        db.drop_all()

We end up very close to the testdriven.io example for unittest.

# project/tests/base.py on http://testdriven.io/part-one-test-setup/
from flask_testing import TestCase
from project import app, db
class BaseTestCase(TestCase):
    def create_app(self):  # done in our pytest fixture before yield
        app.config.from_object('project.config.TestingConfig')
        return app
        
    def setUp(self):  # done in our pytest fixture before yield
        db.create_all()
        db.session.commit()
    def tearDown(self): # done in our pytest fixture after yield
        db.session.remove()
        db.drop_all()

I just haven’t found the reason why db.session.commit() would be necessary (although I suspect it might be an alternative to explicitely declaring the app context like we did in with). I will update the post when I understand it better.