Last week we moved from an in-memory data store to SQLAlchemy and persist our tasks inside an SQLite database. We have two things we need to optimize, or else we end up with problems along the way.
This post is part of my journey to learn Python. You find the code for this post in my PythonFriday repository on GitHub.
The two problems
Our current implementation has two notable problems:
- The tests in test_data_store_db.py write to test_db.sqlite, while those in test_todo.py write to todo_api.sqlite. That can be a problem if the databases differ, and we may accidentally clutter our development database.
- The database(s) will grow with every test run. It may take a long time, but eventually the small entries accumulate and fill our disk.
Let us address the problem with the two databases first, then with a solution for that our second problem is a lot simpler to solve.
Change the get_db() dependency in the endpoints for the tests
The todo_api.sqlite filename comes from our get_db() method in the router. Since we inject the get_db() method into our endpoints as a dependency, we can use the built-in approach of FastAPI to replace dependencies for testing.
In our test, we implement our own get_db() method with the right database file and tell FastAPI to replace the dependency for our tests:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from ..routers.todo import get_db from ..data.datastore_db import DataStoreDb from ..data.database import create_session_factory from ..main import app import logging logging.getLogger("httpx").setLevel(logging.WARNING) def override_get_db(): db_file = os.path.join( os.path.dirname(__file__), '..', 'db', 'test_db.sqlite') factory = create_session_factory(db_file) session = factory() db = DataStoreDb(session) try: yield db finally: session.close() client = TestClient(app) app.dependency_overrides[get_db] = override_get_db |
ATTENTION: The method name for the key in dependency_overrides does not have ‘ ‘ around. If you put them there, the dependency cannot be found and will not be replaced!
Cleanup before the test run
We can clean up the database before or after the test run. The first hunch may be to clean-up after, but that means if something goes wrong, we may no longer have the data to understand the problem. A better approach is to clean up before the test run.
We can define a scope in our with_db() fixture so that it runs only once for the whole test suite. That allows us to delete the old test database (should the file exist) before we run the tests:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
@pytest.fixture(scope="session") def with_db(): db_file = os.path.join( os.path.dirname(__file__), '..', 'db', 'test_db.sqlite') if os.path.isfile(db_file): os.remove(db_file) factory = create_session_factory(db_file) session = factory() try: yield session finally: session.close() |
To make sure that the cleanup does not happen in the middle of our test run, we need to use that fixture everywhere where we use the database. We can extend our endpoint tests to use our fixture:
1 2 |
def test_create_task(with_db): ... |
While this can be done, it is somewhat cumbersome, and we may forget it when we write tests in the future. Let us remove the fixture from the tests we do not need and remove the part that deletes the database file:
1 2 |
# if os.path.isfile(db_file): # os.remove(db_file) |
Startup hook
Pytest lets us use hooks at specific places in the test run. One of the earliest we can use is pytest_sessionstart() that runs after the session object is created, but before the tests are collected.
To use this hook, we create a conftest.py file in the top folder of our project and add this code:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
import os def pytest_sessionstart(session): """ Removes the ./db/test_db.sqlite database at the start of the test run """ db_file = os.path.join( os.path.dirname(__file__), '.', 'db', 'test_db.sqlite') if os.path.isfile(db_file): os.remove(db_file) print(f"DB {db_file} removed") |
Whenever we now run Pytest, the hook runs and if we have an existing test database, it will remove it:
Next
The dependency overwrites and the scope for the test fixture can help us to improve our integration tests. We now only need one database for our tests and every test run starts with a clean, new database.
Next week we integrate Alembic to our application so that we can migrate the database should we need more fields on our entities.