Refactored with much improved stability and less hacks

Author: bharling

README.rst

@@ -41,28 +41,83 @@ Requirements
 Installation
 ------------
 
-You can install "pytest-sqlalchemy" via `pip`_ from `GitHub`_::
+You can install "pytest-sqlalchemy" via `pip`_ from `GitHub`_
+
+.. code-block:: shell
 
     $ pip install -e git+https://github.com/crowdcomms/pytest-sqlalchemy.git#egg=pytest-sqlalchemy
 
 
 Usage
 -----
 
-You need to provide a couple of fixtures to inject your sqlalchemy Base and Session classes ( the ones your regular app uses ). eg::
+The plugin will create a completely new test database that lasts as long as the test session and is destroyed at the end.
+You can execute each test case in an isolated transaction by including the ``db_session`` fixture, eg:
+
+.. code-block:: python
+
+    def test_create_foo(db_session):
+        foo = Foo(name="bar")
+        db_session.add(foo)
+        db_session.commit()
+        assert db_session.query(Foo).count()
+
+The transaction is automatically rolled back at the end of the test function giving you a clean slate for the next test.
+
+You need to define a couple of fixtures to be able to use the plugin, this is mostly to 'patch' your existing Session and Base classes to use the testing database. This is probably the most tricky bit of the plugin as sqlalchemy usage in projects can vary somewhat
+
+Required fixtures
+^^^^^^^^^^^^^^^^^
+``sqlalchemy_base_class``
+
+You must set this to your base class that you use for defining models in your project, eg:
+
+.. code-block:: python
 
     # my_project/tests/conftest.py
-    
-    from my_app.db import Base, Session
+
+    from my_project.database import Base
 
     @pytest.fixture(scope='session')
     def sqlalchemy_base():
         return Base
 
+``sqlalchemy_session_class``
+
+Use this fixture to supply your project's sqlalchemy Session factory, eg:
+
+.. code-block:: python
+
+    # my_project/database.py
+
+    from sqlalchemy.ext.declarative import declarative_base
+    from sqlalchemy.orm import scoped_session, sessionmaker
+
+    Base = declarative_base()
+    Session = scoped_session(sessionmaker())
+
+    ...
+
+    # my_project/tests/conftest.py
+
+    from my_project.database import Session
+
     @pytest.fixture(scope='session')
-    def sqlalchemy_session():
+    def sqlalchemy_session_class():
         return Session
 
+If your project uses a different way to obtain a sqlalchemy session, then you'll need to figure out some other way to configure that session to use the test database, possibly by mocking it in individual test cases.
+
+Optional Fixtures
+^^^^^^^^^^^^^^^^^
+
+``database_url``
+
+This defaults to ``os.environ['DATABASE_URL']`` but is designed to be overridden to supply an alternative. The plugin will attempt to connect to whatever database is specified and create another database alongside the original, prefixed with ``test_``
+
+``test_db_prefix``
+
+If you don't like ``test_`` as a prefix for your testing database, return something else here.
 
 Contributing
 ------------

example_app/db.py

@@ -1,22 +1,13 @@
-from sqlalchemy import Column, Integer, String, Table, ForeignKey, DateTime
-from sqlalchemy.orm import relationship
+from sqlalchemy import Column, Integer, String
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy import create_engine
 from sqlalchemy.orm import scoped_session, sessionmaker
-import os
 
-engine = create_engine(os.environ['DATABASE_URL'])
-
-Base = declarative_base(bind=engine)
-
-Session = sessionmaker(expire_on_commit=False, bind=engine)
-
-session = Session()
+Base = declarative_base()
+Session = scoped_session(sessionmaker())
 
 
 class User(Base):
-
     __tablename__ = 'users'
 
     id = Column(Integer, primary_key=True)
-    name =  Column(String)
+    name = Column(String)

pytest_sqlalchemy.py

@@ -1,16 +1,14 @@
 # -*- coding: utf-8 -*-
-import pytest
+import os
+
 from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
-# from db.models import Base, Session, App
-import contextlib
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import ProgrammingError
 import logging
 import pytest
-import sqlalchemy
-import os
-import importlib
 
-log = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
+
 
 def pytest_addoption(parser):
     group = parser.getgroup('sqlalchemy')
@@ -27,8 +25,8 @@ def pytest_addoption(parser):
 
 
 @pytest.fixture(scope='session')
-def db_prefix(request):
-    return request.config.option.test_db_prefix
+def test_db_prefix():
+    return 'test_'
 
 
 @pytest.fixture(scope='session')
@@ -37,146 +35,68 @@ def database_url():
 
 
 @pytest.fixture(scope='session')
-def test_db(request, database_url, db_prefix):
-    """
-    Create a testing database for the test session.
-    Returns an instance of "TestingDB".
-    """
-    test_db = TestingDB(
-        base_db_url=database_url,
-        drop_existing=request.config.getini('drop_existing_test_db'),
-        prefix=db_prefix
-    )
-    test_db.create()
-    request.addfinalizer(test_db.drop)
-    return test_db
+def test_database_url(test_db_prefix, database_url):
+    test_url = make_url(database_url)
+    test_url.database = test_db_prefix + test_url.database
+    return test_url
 
 
 @pytest.fixture(scope='session')
-def test_engine(request, test_db):
-    """
-    A database Engine connected to the testing database.
-    """
-    engine = test_db.create_engine()
-
-    @request.addfinalizer
-    def cleanup():
-        log.info("Disposing test_engine")
-        engine.dispose()
-
-    return engine
-
-
-class TestingDB:
-    """
-    A test database used for a testing session.
-    This class just provides an empty test database which then can be used to
-    create the schema elements.
-    """
-
-    def __init__(self, base_db_url, drop_existing=False, prefix='test'):
-        self._base_db_url = self._to_url(base_db_url)
-        self._db_url = self._create_test_url(prefix)
-        self._drop_existing = drop_existing
-
-    @property
-    def database(self):
-        return self._db_url.database
-
-    @property
-    def url(self):
-        return self._db_url
-
-    def create(self):
-        log.info("Creating test database %s", self.database)
-        with self.connect() as conn:
-            if self._drop_existing:
-                self._try_drop(conn)
-            conn.execute("CREATE DATABASE {}".format(self.database))
-
-    def _try_drop(self, conn):
-        try:
-            conn.execute("DROP DATABASE {}".format(self.database))
-        except sqlalchemy.exc.ProgrammingError:
-            pass
-        finally:
-            conn.execute("ROLLBACK")
-
-    def drop(self):
-        log.info("Dropping test database %s", self.database)
-        with self.connect() as conn:
-            conn.execute("DROP DATABASE {}".format(self.database))
-
-    @contextlib.contextmanager
-    def connect(self):
-        engine = sqlalchemy.create_engine(self._base_db_url)
-        conn = engine.connect()
-        conn.execution_options(autocommit=False)
-        conn.execute("ROLLBACK")
-        yield conn
-        conn.close()
-        engine.dispose()
-
-    def create_engine(self):
-        """
-        Create an Engine for the test database.
-        """
-        log.info("Creating new Engine for %s", self.database)
-        return sqlalchemy.create_engine(
-            self.url,
-            # TODO: johbo: Without this pool I see trouble due to a left
-            # connection which prevents us dropping the database. Needs
-            # investigation how to solve this in a proper way. Based on the
-            # docs I would expect StaticPool to just work fine, but for some
-            # reason it does not.
-            poolclass=sqlalchemy.pool.AssertionPool,
-            echo=False,
-            echo_pool=False,
-            connect_args={'options': '-c timezone=utc'}
-        )
-
-    def _to_url(self, url):
-        # Ensure to create a copy actually
-        url = str(url)
-        return sqlalchemy.engine.url.make_url(url)
-
-    def _create_test_url(self, prefix):
-        test_url = self._to_url(self._base_db_url)
-        test_url.database = prefix + self._base_db_url.database
-        return test_url
+def test_db(database_url, test_database_url):
+    engine = create_engine(database_url)
+    conn = engine.connect()
+    conn.execution_options(autocommit=False)
+    conn.execute('ROLLBACK')
+
+    try:
+        conn.execute(f"DROP DATABASE {test_database_url.database}")
+    except ProgrammingError:
+        pass
+    finally:
+        conn.execute('ROLLBACK')
+
+    logger.debug('Creating Test Database {}'.format(test_database_url.database))
 
+    conn.execute("CREATE DATABASE {}".format(test_database_url.database))
+
+    conn.close()
+    engine.dispose()
 
 
 @pytest.fixture(scope='session')
 def sqlalchemy_base():
-    return None
+    raise ValueError('Please supply sqlalchemy_base fixture')
 
 
 @pytest.fixture(scope='session')
-def sqlalchemy_session():
-    return None
+def sqlalchemy_session_class():
+    raise ValueError('Please supply sqlalchemy_session_class fixture')
 
 
 @pytest.fixture(scope='session')
-def connection(request, test_engine, sqlalchemy_base, sqlalchemy_session):
-
-    sqlalchemy_base.metadata.bind = test_engine
-    sqlalchemy_base.metadata.create_all(test_engine)
+def engine(test_database_url):
+    return create_engine(test_database_url)
 
-    with test_engine.connect() as conn:
-        yield conn
 
-    sqlalchemy_base.metadata.drop_all()
+@pytest.yield_fixture(scope='session')
+def tables(engine, sqlalchemy_base, test_db):
+    sqlalchemy_base.metadata.create_all(engine)
+    yield
+    sqlalchemy_base.metadata.drop_all(engine)
 
 
-@pytest.fixture
-def db_session(request, connection, sqlalchemy_session, monkeypatch):
+@pytest.yield_fixture(scope='function')
+def db_session(engine, tables, sqlalchemy_session_class):
+    sqlalchemy_session_class.remove()
+    with engine.connect() as connection:
+        transaction = connection.begin_nested()
+        sqlalchemy_session_class.configure(bind=connection)
+        session = sqlalchemy_session_class()
 
-    transaction = connection.begin()
-    sqlalchemy_session.bind = connection
+        session.begin_nested()
 
-    yield sqlalchemy_session
+        yield session
 
-    transaction.rollback()
-    sqlalchemy_session.close()
+        session.close()
+        transaction.rollback()
 

tests/conftest.py

@@ -1,12 +1,14 @@
 import pytest
 pytest_plugins = 'pytester'
 
-from example_app.db import Base, session
+from example_app.db import Base, Session
+
 
 @pytest.fixture(scope='session')
 def sqlalchemy_base():
     return Base
 
+
 @pytest.fixture(scope='session')
-def sqlalchemy_session():
-    return session
+def sqlalchemy_session_class():
+    return Session

tests/test_db_fixture.py

@@ -1,11 +1,13 @@
 from example_app.db import User
 
+
 def test_create_user(db_session):
     user = User(name='The User')
     db_session.add(user)
     db_session.flush()
     assert db_session.query(User).count() == 1
 
+
 def test_update_object(db_session):
     user = User(name='The User')
     db_session.add(user)
@@ -17,5 +19,6 @@ def test_update_object(db_session):
 
     db_session.query(User).filter(User.name=='Bongo').one()
 
+
 def test_tests_run_in_transactions(db_session):
-    assert db_session.query(User).count() == 0
+    assert db_session.query(User).count() == 0

tests/test_pytest_sqlalchemy.py

@@ -3,8 +3,8 @@
 def test_db_prefix_fixture(testdir):
 
     testdir.makepyfile("""
-        def test_dbprefix(db_prefix):
-            assert db_prefix == "test"
+        def test_dbprefix(test_db_prefix):
+            assert test_db_prefix == "test_"
     """)
 
     result = testdir.runpytest(