Flask unittest and sqlalchemy using all connections

问题

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

回答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()

回答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

来源：https://stackoverflow.com/questions/18291180/flask-unittest-and-sqlalchemy-using-all-connections