Test-Driven Development with Python, Testcontainers, and pytest

In this blog post, we will explore a step-by-step guide on how to set up a Python Testcontainers project using pytest for testing a simple application that manages customer data. We will cover the installation of necessary dependencies, creating the project structure, writing test cases, and running the tests.

Step 1: Set up a virtual environment

Create a virtual environment to isolate the project dependencies.

python3 -m venv .venv
source .venv/bin/activate

Step 2: Install required packages

Update pip and install the necessary packages for the project.

pip install --upgrade pip
pip install psycopg pytest testcontainers

Step 3: Create the project structure

Create the following directory structure for your project:

project/
├── customers/
│   ├── __init__.py
│   └── customers.py
├── db/
│   ├── __init__.py
│   └── connection.py
├── tests/
│   ├── __init__.py
│   ├── conftest.py
│   ├── test_customers.py
├── .gitignore
├── Makefile
├── README.md
├── requirements.txt
└── setup.py

Step 4: Implement the customer module

In the customers.py file, implement the customer functionality.
from db.connection import get_connection

def create_table():
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS customers (
            id SERIAL PRIMARY KEY,
            name VARCHAR(255) NOT NULL,
            email VARCHAR(255) UNIQUE NOT NULL
        );
    """)
    conn.commit()
    cursor.close()
    conn.close()

def create_customer(name, email):
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("INSERT INTO customers (name, email) VALUES (%s, %s)", (name, email))
    conn.commit()
    cursor.close()
    conn.close()

def get_all_customers():
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM customers")
    customers_list = cursor.fetchall()
    cursor.close()
    conn.close()
    return customers_list

def get_customer_by_email(email):
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM customers WHERE email = %s", (email,))
    customer = cursor.fetchone()
    cursor.close()
    conn.close()
    return customer

def delete_all_customers():
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM customers")
    conn.commit()
    cursor.close()
    conn.close()

Step 5: Implement the database connection module

In the connection.py file, implement the database connection functionality.

import os
import psycopg2
from testcontainers.postgres import PostgresContainer

postgres = PostgresContainer("postgres:16-alpine")

def get_connection():
    if os.getenv("DB_CONN"):
        return psycopg2.connect(os.getenv("DB_CONN"))
    else:
        postgres.start()
        conn = psycopg2.connect(
            host=postgres.get_container_host_ip(),
            port=postgres.get_exposed_port(5432),
            user=postgres.username,
            password=postgres.password,
            dbname=postgres.dbname
        )
        return conn

Step 6: Write test cases

In the test_customers.py file, write test cases using pytest.

import os
import pytest
from testcontainers.postgres import PostgresContainer

from customers import customers

postgres = PostgresContainer("postgres:16-alpine")


@pytest.fixture(scope="module", autouse=True)
def setup(request):
    """
    Setup the test environment

    This fixture will only run once for all the tests in the module.

    :param request:
    :return:
    """
    postgres.start()

    def remove_container():
        postgres.stop()

    request.addfinalizer(remove_container)
    os.environ["DB_CONN"] = postgres.get_connection_url()
    os.environ["DB_HOST"] = postgres.get_container_host_ip()
    os.environ["DB_PORT"] = postgres.get_exposed_port(5432)
    os.environ["DB_USERNAME"] = postgres.username
    os.environ["DB_PASSWORD"] = postgres.password
    os.environ["DB_NAME"] = postgres.dbname
    customers.create_table()


@pytest.fixture(scope="function", autouse=True)
def setup_data():
    customers.delete_all_customers()


def test_get_all_customers():
    customers.create_customer("Siva", "[e-mail]")
    customers.create_customer("James", "[e-mail]")
    customers_list = customers.get_all_customers()
    assert len(customers_list) == 2


def test_get_customer_by_email():
    customers.create_customer("John", "[e-mail]")
    customer = customers.get_customer_by_email("[e-mail]")
    assert customer.name == "John"
    assert customer.email == "[e-mail]"

Step 7: Run the tests

To run the tests, execute the following command in the project directory:

pytest

================================ test session starts ================================
platform darwin -- Python 3.11.6, pytest-8.3.2, pluggy-1.5.0 -- /Users/testuser/testcontainer-python-demo/.venv/bin/python3.11
cachedir: .pytest_cache
rootdir: /Users/ajeetsraina/devrel/14aug/testcontainer-python-demo
configfile: setup.cfg
testpaths: tests
collected 2 items                                                                   

tests/test_customers.py::test_get_all_customers PASSED                        [ 50%]
tests/test_customers.py::test_get_customer_by_email PASSED                    [100%]

================================ 2 passed in 21.41s =================================

Conclusion:

In this blog post, we have walked through a step-by-step guide on setting up a Python Testcontainers project using pytest for testing a simple application that manages customer data. We have covered the installation of necessary dependencies, creating the project structure, writing test cases, and running the tests. This approach demonstrates the power of Test-Driven Development and the ability to isolate dependencies using Docker containers.