Python database unit testing class

db_test.py

import os
import tempfile
import unittest

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from youapp.model import Base   # replace this with your app's model base

"""
HOW TO USE
Add this logic to your tests __init__ class.
Replace your apps model as necessary (yourapp.model)
Use DbTestClass instead of UnitTest.testcase when creating your test classes
Access database connection details via DbTestCase class methods.

Most basic test class:

    from tests import DbTestCase
    class MyDbTestCase(DbTestCase):
        def test_test(self):
            self.session.add(X)
        self.session.commit()

        # test standard query
        query = self.session.query(X).all()
        self.assertEqual(1, len(query))
        
By default will run against sqlite. Set TEST_DB_TYPE to run against different databases.
"""

class SqliteDriver(object):
    """
    This is the default database driver used for first pass unit tests
    This driver does not require any additional components to be installed
    beyond python itself.
    Uses temp files over memory so data can be inspected after run.
    """
    def __init__(self):
        self.db_file_handle, self.db_file_name = tempfile.mkstemp()
        self.db_url = 'sqlite:///{}'.format(self.db_file_name)
        self.engine = create_engine(self.db_url)

    def destroy(self):
        os.close(self.db_file_handle)
        os.remove(self.db_file_name)


class PostgresqlDriver(object):
    """
    This driver relies on a Postgresql database being available
    If you already have postgresql installed, use the following commands
    to setup a test database for use by these tests

       $ createuser --password -U pricing   # enter 'password' when prompted for password
       $ createdb --owner=pricing pricing_test

    Alternatively use one of the official docker images
    https://hub.docker.com/_/postgres

    Make sure you pass the following environment variables:

    * POSTGRES_PASSWORD = password
    * POSTGRES_USER = pricing
    * POSTGRES_DB = pricing_test

    """
    def __init__(self):
        self.db_url = 'postgresql://pricing:password@localhost/pricing_test'
        self.engine = create_engine(self.db_url)

    def destroy(self):
        Base.metadata.drop_all(self.engine)


class OracleDriver(object):
    """
    This driver relies on an oracle db being available.
    You will need to install the oracle python dbapi - cx_Oracle
    The following docker command will make the required database available on your local machine.
    You will need to setup account on the docker hub and agree to the Oracle licencing terms.
    Before pulling the image, you will need to use the docker login command.

        docker run -d -it --name oracle12 -p 1521:1521 -e DB_SID=PRICING -e DB_PASSWD="p@ssw0rd" store/oracle/database-enterprise:12.2.0.1-slim

    """
    def __init__(self):
        self.db_url = 'oracle://system:p@[email protected]/PRICING'
        self.engine = create_engine(self.db_url)

    def destroy(self):
        Base.metadata.drop_all(self.engine)


class SqlServerDriver(object):
    """
    This driver relies on a SQL server instance being available.
    First, you will need to install the Sql server dbapi driver - pymssql
    The following docker commands will make a Linux SQL Server database available on your local machine

        docker run -e "ACCEPT_EULA=Y" -e "SA_PASSWORD=Passw0rd" -p 1433:1433 --name sql1 -d mcr.microsoft.com/mssql/server:2017-latest
        docker exec -it sql1 /opt/mssql-tools/bin/sqlcmd -U sa -P Passw0rd -Q "CREATE DATABASE PRICING_TEST;"

    """
    def __init__(self):
        self.db_url = 'mssql+pymssql://sa:Passw0rd@localhost/PRICING_TEST?charset=utf8'
        self.engine = create_engine(self.db_url)

    def destroy(self):
        Base.metadata.drop_all(self.engine)


class DbNoSetupTestCase(unittest.TestCase):
    def get_db_driver(self):
        db_type = os.environ.get('TEST_DB_TYPE')
        if db_type == 'postgresql':
            return PostgresqlDriver()
        if db_type == 'oracle':
            return OracleDriver()
        if db_type == 'mssql':
            return SqlServerDriver()
        return SqliteDriver()

    def setUp(self):
        # load the db driver object
        self.db_driver = self.get_db_driver()
        self.engine = self.db_driver.engine
        self.db_url = self.db_driver.db_url

    def tearDown(self):
        self.db_driver.destroy()


class DbTestCase(DbNoSetupTestCase):
    def setUp(self):
        super(DbTestCase, self).setUp()
        Base.metadata.create_all(self.engine)
        self.session = sessionmaker(bind=self.engine)()

    def tearDown(self):
        self.session.commit()
        self.session.close()
        super(DbTestCase, self).tearDown()