修改Flask的SQLALCHEMY_DATABASE_URI后测试仍连接主数据库问题

问题描述 #

测试基于Flask和SQLAlchemy的待办事项应用时，在测试文件中修改app.config['SQLALCHEMY_DATABASE_URI']为任意值后，虽然打印确认配置已变更，但数据库提交操作仍在main.py指定的主数据库执行。尝试过通过os.environ修改URI、使用Flask-Testing包、在导入语句前修改配置等方法，均无法解决问题。

原代码 #

main.py #

from flask import Flask, render_template, request, redirect, url_for
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///db.sqlite'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)

class Todo(db.Model):

    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(100))
    complete = db.Column(db.Boolean)

@app.route("/")
def home():
    todo_list = Todo.query.all()
    return render_template("base.html", todo_list=todo_list)

@app.route("/add", methods=["POST"])
def add():
    title = request.form.get("title")
    new_todo = Todo(title=title, complete=False)
    db.session.add(new_todo)
    db.session.commit()
    return redirect(url_for("home"))

@app.route("/update/<int:todo_id>")
def update(todo_id):
    todo = Todo.query.filter_by(id=todo_id).first()
    todo.complete = not todo.complete
    db.session.commit()
    return redirect(url_for("home"))

@app.route("/delete/<int:todo_id>")
def delete(todo_id):
    todo = Todo.query.filter_by(id=todo_id).first()
    db.session.delete(todo)
    db.session.commit()
    return redirect(url_for("home"))

if __name__ == "__main__":
    with app.app_context():
        db.create_all()
        app.run(debug=True)

test_main.py #

from main import db, app, Todo
import unittest

class TodoTests(unittest.TestCase):

    def setUp(self):
        app.config["TESTING"] = True
        app.config["WTF_CSRF_ENABLED"] = False
        app.config["DEBUG"] = False
        app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
        self.app = app.test_client()
        db.drop_all()
        db.create_all()
        self.create_todo()
        self.assertEqual(app.debug, False)

    def create_todo(self):
        item = Todo(title="Todo")
        db.session.add(item)
        db.session.commit()

    def test_todo_endpoint(self):
        response = self.app.get("/")
        # print (response.data)
        self.assertIn(
            b'Todo', response.data
        )
        self.assertEqual(response.status_code, 200)

if __name__ == "__main__":
    unittest.main()

问题根源 #

SQLAlchemy实例在导入main模块时就已与主数据库绑定，后续修改app配置不会更新已建立的数据库连接。

当测试文件执行from main import db, app, Todo时，main.py中的db = SQLAlchemy(app)已经执行完毕，此时db实例已和配置了主数据库URI的app绑定并建立连接池。后续在setUp中修改app.config['SQLALCHEMY_DATABASE_URI']仅会改变app的配置值，但不会触发db实例重新初始化连接，因此所有数据库操作仍指向主数据库。

解决方案 #

1. 重构main.py为应用工厂模式（推荐） #

应用工厂模式能灵活创建不同配置的app实例，完美适配测试场景：

修改后的main.py：

from flask import Flask, render_template, request, redirect, url_for
from flask_sqlalchemy import SQLAlchemy

# 先初始化SQLAlchemy实例，暂不绑定app
db = SQLAlchemy()

class Todo(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(100))
    complete = db.Column(db.Boolean)

def create_app(config=None):
    app = Flask(__name__)
    # 默认配置
    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///db.sqlite'
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    
    # 传入自定义配置时覆盖默认值
    if config:
        app.config.update(config)
    
    # 将app与SQLAlchemy实例绑定
    db.init_app(app)
    
    # 注册路由
    @app.route("/")
    def home():
        todo_list = Todo.query.all()
        return render_template("base.html", todo_list=todo_list)

    @app.route("/add", methods=["POST"])
    def add():
        title = request.form.get("title")
        new_todo = Todo(title=title, complete=False)
        db.session.add(new_todo)
        db.session.commit()
        return redirect(url_for("home"))

    @app.route("/update/<int:todo_id>")
    def update(todo_id):
        todo = Todo.query.filter_by(id=todo_id).first()
        todo.complete = not todo.complete
        db.session.commit()
        return redirect(url_for("home"))

    @app.route("/delete/<int:todo_id>")
    def delete(todo_id):
        todo = Todo.query.filter_by(id=todo_id).first()
        db.session.delete(todo)
        db.session.commit()
        return redirect(url_for("home"))
    
    return app

# 创建默认应用实例
app = create_app()

if __name__ == "__main__":
    with app.app_context():
        db.create_all()
        app.run(debug=True)

2. 更新测试文件test_main.py #

通过应用工厂创建带测试配置的app，确保使用内存数据库：

from main import db, create_app, Todo
import unittest

class TodoTests(unittest.TestCase):
    def setUp(self):
        # 创建测试专用app，传入自定义配置
        self.app = create_app({
            "TESTING": True,
            "WTF_CSRF_ENABLED": False,
            "DEBUG": False,
            "SQLALCHEMY_DATABASE_URI": "sqlite:///:memory:"
        })
        self.client = self.app.test_client()
        
        # 在测试app的上下文内初始化数据库
        with self.app.app_context():
            db.create_all()
            self.create_todo()

    def create_todo(self):
        item = Todo(title="Todo")
        db.session.add(item)
        db.session.commit()

    def test_todo_endpoint(self):
        response = self.client.get("/")
        self.assertIn(b'Todo', response.data)
        self.assertEqual(response.status_code, 200)
    
    def tearDown(self):
        # 测试结束后清理数据库
        with self.app.app_context():
            db.session.remove()
            db.drop_all()

if __name__ == "__main__":
    unittest.main()

额外说明 #

• 使用sqlite:///:memory:时，必须确保所有数据库操作都在同一个app上下文内，否则会创建独立的内存数据库，导致测试数据无法共享。
• 若不想重构为应用工厂，也可以在测试的setUp中修改配置后，重新初始化db实例，但这种方式耦合性高，不推荐。

内容的提问来源于stack exchange，提问作者Kate