Demonstrate flush / rollback to handle database state

main.py

from flask import Flask, jsonify
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)
app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://test@localhost/test'
db = SQLAlchemy(app)
db.init_app(app)


class User(db.Model):
    """
    Test users
    """
    __tablename__ = 'users'

    user_id = db.Column(db.String, primary_key=True)

    def to_json(self):
    	return dict(user_id=self.user_id)

db.create_all()

@app.route('/users', methods=['GET'])
def users():
    users = db.session.query(User)
    output = {'users': [user.to_json() for user in users]}
    return jsonify(**output)

pseudotest.py

from main import app, User, db
import json

ctx = app.app_context()

ctx.push()

def test_users():
	user1 = User(user_id='one')
	user2 = User(user_id='two')
	db.session.add_all([user1, user2])
	db.session.flush()
	assert User.query.count() > 0
	response = app.test_client().get('/users')
	assert response.status.startswith('200')
	assert json.loads(response.data) == {u'users': [{u'user_id': u'one'}, {u'user_id': u'two'}]}
	db.session.rollback()
	assert User.query.count() == 0


if __name__ == '__main__':
	test_users()

setup_and_run.sh

#!/bin/bash
createdb test
createuser test
python pseudotest.py