Flask unittest and sqlalchemy using all connections - postgresql

Question

I've just run into an issue running unittests on my flask app after I had roughly 100 unittests. All unittests will pass, but when run all at once they will fail with the following error:
OperationalError: (OperationalError) FATAL: remaining connection slots are reserved for non-replication superuser connections
Everything is running in a virtualbox/vagrant/ubuntu12.04 instance on local machine. My postgres max_connections is set to 100 so I'm assuming that the connections aren't closing and after running 100 tests I use up all the available ones.
This person Flask unit tests with SQLAlchemy and PostgreSQL exhausts db connections looks like they are having the same exact problem. Mike/Zzzeek (sqlalchemy dev) even responded to it saying that something may be happening in create_app() so I've included that as well below.
Does this mean I'm not closing my connections somewhere? All of these errors are triggered by db.create_all() in my setUp() method of my unittest.
# test.py
class TestCase(DataMixin, Base):
"""Base test class"""
def create_app(self):
return create_app(TestConfig())
def setUp(self):
db.create_all()
def tearDown(self):
db.session.remove()
db.drop_all()
# app.py
def create_app(config=None):
app = Flask(__name__)
# Config
app.config.from_object(BaseConfig())
if config is not None:
app.config.from_object(config)
# Extensions
db.init_app(app)
mail.init_app(app)
bcrypt.init_app(app)
# Blueprints
app.register_blueprint(core_blueprint, url_prefix='/')
app.register_blueprint(accounts_blueprint, url_prefix='/account')
app.register_blueprint(admin_blueprint, url_prefix='/admin')
app.register_blueprint(cart_blueprint, url_prefix='/cart')
# Login Manager
login_manager.setup_app(app, add_context_processor=True)
login_manager.login_view = "accounts.login"
login_manager.user_callback = load_user
# Templates
app.jinja_env.globals['is_admin'] = is_admin
app.jinja_env.globals['is_staff'] = is_staff
#app.context_processor
def inject_cart():
cart = count = None
if current_user.is_authenticated():
cart = current_user.get_cart()
return dict(cart=cart)
# Error Handling
#app.errorhandler(404)
def page_not_found(error):
return render_template('404.html'), 404
return app

Answer 1

UPDATE: Tested and fixed
Instead of making a new connection and re-creating your database every time (slow), you can use subsessions and do a rollback after each test.
The connection are reused, so this also fix the problem you're having.
class TestCase(Base):
#classmethod
def setUpClass(cls):
cls.app = create_app(MyConfig())
cls.client = cls.app.test_client()
cls._ctx = cls.app.test_request_context()
cls._ctx.push()
db.create_all()
#classmethod
def tearDownClass(cls):
db.session.remove()
db.drop_all()
db.get_engine(cls.app).dispose()
def setUp(self):
self._ctx = self.app.test_request_context()
self._ctx.push()
db.session.begin(subtransactions=True)
def tearDown(self):
db.session.rollback()
db.session.close()
self._ctx.pop()
If you need to also make an instance of the application for each test, just add it to the setUp method but leave it also in setUpClass.
Full test example below requires flask_sqlalchemy and psycopg2. Create a test database named "test" and set its connection limit to 15.
from flask import Flask
from flask.ext.sqlalchemy import SQLAlchemy
from unittest import TestCase as Base
db = SQLAlchemy()
def create_app(config=None):
app = Flask(__name__)
app.config.from_object(config)
db.init_app(app)
return app
class MyConfig(object):
SQLALCHEMY_DATABASE_URI = "postgresql://localhost/test"
TESTING = True
class TestCase(Base):
#classmethod
def setUpClass(cls):
cls.app = create_app(MyConfig())
cls.client = cls.app.test_client()
cls._ctx = cls.app.test_request_context()
cls._ctx.push()
db.create_all()
#classmethod
def tearDownClass(cls):
db.session.remove()
db.drop_all()
def setUp(self):
self._ctx = self.app.test_request_context()
self._ctx.push()
db.session.begin(subtransactions=True)
def tearDown(self):
db.session.rollback()
db.session.close()
self._ctx.pop()
class TestModel(TestCase):
def test_01(self):
pass
def test_02(self):
pass
def test_03(self):
pass
def test_04(self):
pass
def test_05(self):
pass
def test_06(self):
pass
def test_07(self):
pass
def test_08(self):
pass
def test_09(self):
pass
def test_10(self):
pass
def test_11(self):
pass
def test_12(self):
pass
def test_13(self):
pass
def test_14(self):
pass
def test_15(self):
pass
def test_16(self):
pass
if __name__ == "__main__":
import unittest
unittest.main()

Answer 2

I found the answer here -- https://stackoverflow.com/a/17998485/1870623 and a great explination here -- https://stackoverflow.com/a/16390645/1870623
The solution is to add db.get_engine(self.app).dispose() to the tearDown()
class TestCase(Base):
def setUp(self):
db.create_all()
def tearDown(self):
db.session.remove()
db.drop_all()
db.get_engine(self.app).dispose() # This