
from flask.ext.sqlalchemy import SessionBase

class TestingSession(SessionBase):
    def __init__(self, db, bind, **options):
        self.app = db.get_app()

        SessionBase.__init__(
            self, autocommit=False, autoflush=True,
            bind=bind, **options
        )

    def __call__(self):
        # Flask-SQLAlchemy wants to create a new session
        # Simply return the existing session
        return self

    def get_bind(self, mapper=None, clause=None):
        # mapper is None if someone tries to just get a connection
        if mapper is not None:
            info = getattr(mapper.mapped_table, 'info', {})
            bind_key = info.get('bind_key')
            if bind_key is not None:
                state = flask.ext.sqlalchemy.get_state(self.app)
                return state.db.get_engine(self.app, bind=bind_key)
        return SessionBase.get_bind(self, mapper, clause)

from unittest import TestCase

from app import db

def setUpModule():
    """Sets up the module so that we can run the whole test suite inside
    a single transaction. Each test will be executed in a nested transaction
    using PostgreSQL savepoints.

    We will:

    1. Start a transaction
    2. Drop everything
    3. (Re-)create everything
    """
    global session_backup, connection

    # Create a connection and start a transaction. This is needed so that
    # we can run the drop_all/create_all inside the same transaction as
    # the tests
    connection = db.engine.connect()
    transaction = connection.begin()

    # Back up the original session and replace with our own
    session_backup = db.session
    #db.session = db.sessionmaker()(bind=connection)
    db.session = TestingSession(db, connection)

    ## Drop all to get an empty database free of old crud
    db.metadata.drop_all(transaction.connection)

    ## Create everything
    db.metadata.create_all(transaction.connection)



def tearDownModule():
    """Roll back everything, leaving database as it was before we started"""
    db.session.rollback()
    connection.close()

    # Restore backup
    db.session = session_backup


class FlaskDbTestCase(TestCase):
    @classmethod
    def setUpClass(cls):
        """Makes a savepoint for this class"""
        cls.savepoint1 = db.session.begin_nested()
        cls.savepoint2 = db.session.begin_nested()

        # Replace app db object temporarily with a nested tx
        cls.session_backup = db.session
        db.session = cls.savepoint2.session

    @classmethod
    def tearDownClass(cls):
        """Rolls back to the class savepoint"""
        cls.savepoint1.rollback()

        # Restore original session
        db.session = cls.session_backup

    def setUp(self):
        """Makes a savepoint for this test"""
        # We actually need to make 2 savepoints in case the code to be tested
        # happens to issue a commit so that we can roll back to this point.
        # A commit in the tested code can only commit "savepoint2" so we
        # will use "savepoint" for the rollback.
        self.savepoint1 = db.session.begin_nested()
        self.savepoint2 = db.session.begin_nested()

        # Replace app db object temporarily with a nested tx
        self.session_backup = db.session
        db.session = self.savepoint2.session

        # This is for using app_context().pop()
        db.session.remove = lambda: None

    def tearDown(self):
        """Rolls back to the test class savepoint"""
        # Roll back savepoint - schema will back to fresh
        self.savepoint1.rollback()

        # Restore original session
        db.session = self.session_backup



