Mocking database tables with the Python peewee ORM is super simple.

Of all the Python ORMs I’ve tried, Peewee is my favourite. It is simple, get the jobs done without too much abstraction. One of the benefits is that it is super easy to mock when running unit tests.

For production, Postgres is my go-to database, but for testing, it is much easier to use SQLite. Thankfully making the switch is easy.

Here is an example of a unit test:

@with_test_db((MyTable,))def test_task():     MyTable.create(tenant_id='TEST', name="Bob")     # Run code that uses MyTableOR@with_test_db((MyTable, MyTable2))def test_task():     MyTable.create(tenant_id='TEST', name="Bob")     # Run code that uses MyTable and MyTable2

The magic sauce is in the @with_test_db decorator. This takes care of creating the table and destroying it at the end of the test. The decorator is mostly taken from a blog post by dvlv [1]

from functools import wrapsimport tempfilefrom peewee import SqliteDatabasefrom database_models import Integrationdef with_test_db(dbs: tuple):    """    Decorator to create a test database in memory and bind it to the models.    This will create the tables, run the test function, and then drop the tables and close the database.    :param dbs: A tuple of the models to bind to the test database    """    def decorator(func):        @wraps(func)        def test_db_closure(*args, **kwargs):            test_db = SqliteDatabase(os.path.join(tempfile.gettempdir(), 'test.db'))            # OR            # test_db = SqliteDatabase(":memory:") # works fine for local testing but not with multithreading            with test_db.bind_ctx(dbs):                test_db.create_tables(dbs)                try:                    func(*args, **kwargs)                finally:                    test_db.drop_tables(dbs)                    test_db.close()        return test_db_closure    return decorator

This decorator adds the following steps shown in orange.

One of the challenges with using Peewee however is the way the database connection is used. The database connection is on a per-table basis. The database connection is typically started in the module where the models (tables) are defined.

For unit testing, we don’t want to use Postgres at all so we need to avoid any references to Postgres when in the testing environment. Here we use an environment variable to either start Postgres or create a dummy SQLite database. This database won’t be used during unit testing but it does need to be initialized. An in-memory SQLite database is used for this purpose.

database_models.py:

if os.getenv('ENVIRONMENT') == 'test':    from playhouse.sqlite_ext import *    db = SqliteExtDatabase(':memory')else:    from playhouse.postgres_ext import *    db = PostgresqlExtDatabase(os.getenv('DATABASE'))class BaseModel(Model):    class Meta:        database = db        legacy_table_names = Falseclass MyTable(BaseModel):    version = TextField()

Notes

There are a couple of important things to note.

  1. Stick to SQL field types that exist for both SQLite and Postgres.
  2. Using an SQLite database file is preferable to using a SQLite memory database because, while this works for some cases, it does not work in a multi-threaded or multi-process environment.

References

[1] https://www.dvlv.co.uk/a-super-helpful-decorator-for-peeweeflask-unit-testing.html