В Pytest, всеми любимой среде тестирования Python, приспособление — это многократно используемый фрагмент кода, который упорядочивает что-то перед входом теста и очищает после его завершения. Например, временный файл или папка, среда настройки, запуск веб-сервера и т. д. В этом посте мы рассмотрим , как создать приспособление Pytest, которое создает тестовую базу данных (пустую или с известным состоянием), которая получает очищена, что позволяет запускать каждый тест в полностью чистой базе данных.

Цели

Мы создадим приспособление Pytest с использованием Psycopg 3 для подготовки и очистки тестовой базы данных. Поскольку пустая база данных редко бывает полезна для тестирования, мы дополнительно применим миграцию Yoyo (на момент написания статьи веб-сайт не работает, перейдите к снимку archive.org), чтобы заполнить ее.

Итак, требования к приспособлению Pytest с именем test_db, созданному в этом посте:

• удалить тестовую базу данных, если она существовала до теста
• создайте пустую базу данных перед тестом
• необязательно примените миграцию или создайте тестовые данные перед тестом
• обеспечить подключение к тестовой базе данных к тесту
• удалить тестовую базу данных после теста (даже в случае неудачи)

Любой тестовый метод, который запрашивает его, указав его аргумент тестового метода:

def test_create_admin_table(test_db):
    ...

Получит обычный экземпляр Psycopg Connection, подключенный к тестовой базе данных. Test может делать все, что ему нужно, как и при обычном использовании Psycopg, например:

def test_create_admin_table(test_db):
    # Open a cursor to perform database operations
    cur = test_db.cursor()

    # Pass data to fill a query placeholders and let Psycopg perform
    # the correct conversion (no SQL injections!)
    cur.execute(
        "INSERT INTO test (num, data) VALUES (%s, %s)",
        (100, "abc'def"))

    # Query the database and obtain data as Python objects.
    cur.execute("SELECT * FROM test")
    cur.fetchone()
    # will return (1, 100, "abc'def")

    # You can use `cur.fetchmany()`, `cur.fetchall()` to return a list
    # of several records, or even iterate on the cursor
    for record in cur:
        print(record)

Макет файла проекта

В классическом проекте Python исходные коды находятся в src/, а тесты — в test/:

├── src
│   └── tuvok
│       ├── __init__.py
│       └── sales
│           └── new_user.py
├── tests
│   ├── conftest.py
│   └── sales
│       └── test_new_user.py
├── requirements.txt
└── yoyo.ini

Если вы используете библиотеку миграций, такую ​​​​как фантастическая Yoyo, сценарии миграции, скорее всего, находятся в миграции/:

├── migrations
    ├── 20240816_01_Yn3Ca-sales-user-user-add-last-run-table.py
    ├── ...

Конфигурация

Нашей тестовой базе данных потребуется очень небольшая настройка:

• URL-адрес подключения - (без базы данных)
• имя тестовой базы данных — будет создаваться заново для каждого теста
• (необязательно) папка миграции — сценарии миграции, которые будут применяться для каждого теста

Pytest имеет естественное место conftest.py для совместного использования фикстур в нескольких файлах. Конфигурация прибора тоже будет там:

# Without DB name!
TEST_DB_URL = "postgresql://localhost"
TEST_DB_NAME = "test_tuvok"
TEST_DB_MIGRATIONS_DIR = str(Path(__file__, "../../migrations").resolve())

Вы можете установить эти значения из переменной среды или чего-то еще, подходящего для вашего случая.

Создать приспособление test_db

Зная библиотеку PostgreSQL и Psycopg, напишите фикстуру в conftest.py:

@pytest.fixture
def test_db():
    # autocommit=True start no transaction because CREATE/DROP DATABASE
    # cannot be executed in a transaction block.
    with psycopg.connect(TEST_DB_URL, autocommit=True) as conn:
        cur = conn.cursor()

        # create test DB, drop before
        cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB_NAME}" WITH (FORCE)')
        cur.execute(f'CREATE DATABASE "{TEST_DB_NAME}"')

        # Return (a new) connection to just created test DB
        # Unfortunately, you cannot directly change the database for an existing Psycopg connection. Once a connection is established to a specific database, it's tied to that database.
        with psycopg.connect(TEST_DB_URL, dbname=TEST_DB_NAME) as conn:
            yield conn

        cur.execute(f'DROP DATABASE IF EXISTS "{TEST_DB_NAME}" WITH (FORCE)')

Создать приспособление миграции

В нашем случае мы используем миграцию Yoyo. Напишите применить миграцию как еще одно приспособление под названием yoyo:

@pytest.fixture
def yoyo():
    # Yoyo expect `driver://user:pass@host:port/database_name?param=value`.
    # In passed URL we need to
    url = (
        urlparse(TEST_DB_URL)
        .
        # 1) Change driver (schema part) with `postgresql psycopg` to use
        # psycopg 3 (not 2 which is `postgresql psycopg2`)
        _replace(scheme="postgresql psycopg")
        .
        # 2) Change database to test db (in which migrations will apply)
        _replace(path=TEST_DB_NAME)
        .geturl()
    )

    backend = get_backend(url)
    migrations = read_migrations(TEST_DB_MIGRATIONS_DIR)

    if len(migrations) == 0:
        raise ValueError(f"No Yoyo migrations found in '{TEST_DB_MIGRATIONS_DIR}'")

    with backend.lock():
        backend.apply_migrations(backend.to_apply(migrations))

Если вы хотите применить миграцию к каждой тестовой базе данных, потребуется приспособление yoyo для приспособления test_db:

@pytest.fixture
def test_db(yoyo):
    ...

Чтобы применить миграцию только к некоторым тестам, потребуется yoyo индивидуально:

def test_create_admin_table(test_db, yoyo):
    ...

Заключение

Создание собственного приспособления для предоставления вашим тестам чистой базы данных было для меня полезным опытом, позволившим мне глубже углубиться как в Pytest, так и в Postgres.

Надеюсь, эта статья помогла вам создать собственный набор тестов для баз данных. Не стесняйтесь оставлять мне свой вопрос в комментариях и удачного вам программирования!