Unit testing a function that depends on database

Question

I am running tests on some functions. I have a function that uses database queries. So, I have gone through the blogs and docs that say we have to make an in memory or test database to use such functions. Below is my function,

def already_exists(story_data,c):
    # TODO(salmanhaseeb): Implement de-dupe functionality by checking if it already
    # exists in the DB.
    c.execute("""SELECT COUNT(*) from posts where post_id = ?""", (story_data.post_id,))
    (number_of_rows,)=c.fetchone()
    if number_of_rows > 0:
        return True
    return False

This function hits the production database. My question is that, when in testing, I create an in memory database and populate my values there, I will be querying that database (test DB). But I want to test my already_exists() function, after calling my already_exists function from test, my production db will be hit. How do I make my test DB hit while testing this function?

Answer 1

There are two routes you can take to address this problem:

1. Make an integration test instead of a unit test and just use a copy of the real database.
2. Provide a fake to the method instead of actual connection object.

Which one you should do depends on what you're trying to achieve.

If you want to test that the query itself works, then you should use an integration test. Full stop. The only way to make sure the query as intended is to run it with test data already in a copy of the database. Running it against a different database technology (e.g., running against SQLite when your production database in PostgreSQL) will not ensure that it works in production. Needing a copy of the database means you will need some automated deployment process for it that can be easily invoked against a separate database. You should have such an automated process, anyway, as it helps ensure that your deployments across environments are consistent, allows you to test them prior to release, and "documents" the process of upgrading the database. Standard solutions to this are migration tools written in your programming language like albemic or tools to execute raw SQL like yoyo or Flyway. You would need to invoke the deployment and fill it with test data prior to running the test, then run the test and assert the output you expect to be returned.

If you want to test the code around the query and not the query itself, then you can use a fake for the connection object. The most common solution to this is a mock. Mocks provide stand ins that can be configured to accept the function calls and inputs and return some output in place of the real object. This would allow you to test that the logic of the method works correctly, assuming that the query returns the results you expect. For your method, such a test might look something like this:

from unittest.mock import Mock

...

def test_already_exists_returns_true_for_positive_count():
    mockConn = Mock(
        execute=Mock(),
        fetchone=Mock(return_value=(5,)),
    )
    story = Story(post_id=10) # Making some assumptions about what your object might look like.

    result = already_exists(story, mockConn)

    assert result

    # Possibly assert calls on the mock. Value of these asserts is debatable.
    mockConn.execute.assert_called("""SELECT COUNT(*) from posts where post_id = ?""", (story.post_id,))
    mockConn.fetchone.assert_called()

Answer 2

The issue is ensuring that your code consistently uses the same database connection. Then you can set it once to whatever is appropriate for the current environment.

Rather than passing the database connection around from method to method, it might make more sense to make it a singleton.

def already_exists(story_data):
    # Here `connection` is a singleton which returns the database connection.
    connection.execute("""SELECT COUNT(*) from posts where post_id = ?""", (story_data.post_id,))
    (number_of_rows,) = connection.fetchone()
    if number_of_rows > 0:
        return True
    return False

Or make connection a method on each class and turn already_exists into a method. It should probably be a method regardless.

def already_exists(self):
    # Here the connection is associated with the object.
    self.connection.execute("""SELECT COUNT(*) from posts where post_id = ?""", (self.post_id,))
    (number_of_rows,) = self.connection.fetchone()
    if number_of_rows > 0:
        return True
    return False

But really you shouldn't be rolling this code yourself. Instead you should use an ORM such as SQLAlchemy which takes care of basic queries and connection management like this for you. It has a single connection, the "session".

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from sqlalchemy_declarative import Address, Base, Person

engine = create_engine('sqlite:///sqlalchemy_example.db')
Base.metadata.bind = engine

DBSession = sessionmaker(bind=engine)
session = DBSession()

Then you use that to make queries. For example, it has an exists method.

session.query(Post.id).filter(q.exists()).scalar()

Using an ORM will greatly simplify your code. Here's a short tutorial for the basics, and a longer and more complete tutorial.