Overview

When I first started learning how to use a database, I wrote code that was hard to work with.

To test that functionality worked correctly, I did a lot of local set up that required manual actions like creating tables, adding entries, deleting entries, etc. Sometimes I would deploy my code to see if everything worked properly against the real database.

In this post, I’ll go over how you could organize your code to isolate functionality and make working with a database easier.

Sample Todo App

Here’s a contrived todo application that checks if any todos exist and creates one if there aren’t any.

import os
import psycopg2

db = psycopg2.connect(**{
    'dbname': 'todo_app',
    'user': os.environ.get('DB_USER'),
    'password': os.environ.get('DB_PASSWORD'),
    'port': '54320',
    'host': 'localhost'
})

def generate_todo():
    # get todos from database
    with db.cursor() as cur:
        cur.execute('SELECT todo_id, description FROM todo;')
        todos = cur.fetchall()

    # if there are no todos, create a todo to add more todos
    if not todos:
        with db.cursor() as cur:
            cur.execute('INSERT INTO todo (description) VALUES (%s)', ('Add todos',))
        db.commit()

The problem here is that you can’t verify the functionality without connecting to a real database. You have to run the code and then connect to the database to see if a row was added. You also have to delete the row you added.

The other issue with this code is that it can be difficult to follow. We have database operations and application logic all in one place.

How can we organize this better?

Create a class to simplify database operations and provide it as a dependency

We could define a class that is in charge of all database operations. It provides a simplified interface that can be used by our application.

class DatabaseManager:
    def __init__(self, db):
        self.db = db

    def get_todos(self):
        with self.db.cursor() as cur:
            cur.execute('SELECT todo_id, description FROM todo;')
            results = cur.fetchall()
        return results

    def create_todo(self, description):
        with self.db.cursor() as cur:
            cur.execute('INSERT INTO todo (description) VALUES (%s)', (description,))
            self.db.commit()

Our generate_todo function can now be greatly simplified. It’s also a lot more clear what’s happening. Another benefit is that if we wanted to make changes to db_manager (e.g. use a different database, change the queries, etc) we don’t have to make any changes to generate_todo.

def generate_todo(db_manager):
    todos = db_manager.get_todos()
    if not todos:
        print('Inserting todo')
        db_manager.create_todo('Add todos')

It’s also easy to unit test our new functionality. We don’t need to connect to an actual database. We only need to define a class that fulfills the functionality of the class we defined. Here is a simple unit test using mocks for the db manager.

import unittest
from unittest.mock import Mock
from todo_v2 import generate_todo

class TestGenerateTodo(unittest.TestCase):

    def test_if_no_todos_creates_todo(self):
        db_manager = Mock()
        db_manager.get_todos.return_value = []

        generate_todo(db_manager)

        db_manager.get_todos.assert_called()
        db_manager.insert_todo.assert_called_with('Add todos')

    def test_if_todos_does_not_creates_todo(self):
        db_manager = Mock()
        db_manager.get_todos.return_value = [(1, 'todo')]

        generate_todo(db_manager)

        db_manager.get_todos.assert_called()
        db_manager.insert_todo.assert_not_called()

Testing actual database operations

Sometimes we still want to test real database queries. The benefit of creating a dedicated class for handling database operations is that the functionality is a lot more focused and we can test things in a more isolated manner i.e we don’t have to worry about the application logic.

import unittest
from todo_v2 import DatabaseManager
import psycopg2

class TestDatabaseManager(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.db = psycopg2.connect(**{
            'dbname': 'todo_app',
            'user': 'postgres',
            'password': 'password',
            'port': '54320',
            'host': 'localhost'
        })
        cls.db_manager = DatabaseManager(cls.db)

    def setUp(self):
        with self.db.cursor() as cur:
            cur.execute('TRUNCATE todo')
        self.db.commit()

    def test_create_and_get_todos(self):
        self.db_manager.create_todo('todo')
        todos = self.db_manager.get_todos()
        assert todos[0][1] == 'todo'

if __name__ == '__main__':
    unittest.main()

In this case, you’d still need an actual database to run your tests against. One thing to be careful about is that you don’t want different tests to interfere with each other (notice we clear out the table between tests in setUp). Another problem is if tests were to be run in parallel.

Conclusion

This is an example of Dependency Inversion and it’s a technique I’ve found very helpful in managing my code.

Link to project code: https://github.com/levelupSE/todoapp