この記事は、ソフトウェアテスト Advent Calendar 2023 の13日目の記事です。

ソフトウェアテスト Advent Calendar 2023
https://qiita.com/advent-calendar/2023/softwaretesting

このエントリでは、
pytestのfixtureを使ったPostgreSQLに依存するテストコードの書き方を示します。

DBに対してクエリを実行する処理が含まれるコードは、
単体テストの実行環境を準備するのが面倒になりがちですが、
このエントリのように、fixtureにまとめてしまえば簡単に用意できるので便利です。

また、テストを動かすための、
プロセス実行からDB・ユーザの用意・初期化などもコードで管理できていると、
テスト環境を再現するための資源が一カ所にまとまって管理しやすくなると考えられます。

今回はPostgreSQLを例に示しましたが、
他のDBMS、ミドルウェアも同様の方法でテストができると思います。

PostgreSQLのインストール

まずは、PostgreSQLをインストールします。次のコマンドはUbuntu20.04の場合の例です。

sudo apt install postgresql

Ubuntu20.04の場合、/usr/lib/postgresql/16/bin/postgres のようなパスにインストールされます。

pytest, psycopg2のインストール

pytestとpysycopgをインストールします。

pip install pytest
pip install psycopg2

テストコードの実装

ちょっと長いですが、テストコードは次のように実装します。
大部分はPostgreSQLの起動・終了・DB作成を行うための処理で、
テストの実体は「test_postgresql_query」メソッドのみです。

test/test_postgresql.py

import pytest
import subprocess
import tempfile
import os
import time
import psycopg2

class PostgresqlService():
    def __init__(
            self,
            user="postgres_user", password="postgres_password", db="postgres_db", host="127.0.0.1", port=5432,
            health_check_repeat=5, health_check_wait=1
        ):
        self.user = user
        self.password = password
        self.db = db
        self.host = host
        self.port = port
        self.health_check_repeat = health_check_repeat
        self.health_check_wait = health_check_wait
        self.dsn = f"postgres://{user}:{password}@{host}:{port}/{db}"
        self.pre_dsn = f"postgres://{user}:{password}@{host}:{port}"
        self.working_dir = tempfile.TemporaryDirectory()
        self.data_dir = os.path.join(self.working_dir.name, "data")
        self.socket_dir = self.working_dir.name
        pgbinroot = os.getenv("PGBINROOT", "/usr/lib/postgresql")
        latest_version = sorted([
            int(e)
            for e in os.listdir(pgbinroot)
            if e.isdecimal() and e.isalnum()
        ])[-1]
        self.postgres_bin_dir = os.path.join(pgbinroot, str(latest_version), "bin")
        self.args = [
            os.path.join(self.postgres_bin_dir, "postgres"),
            "-D", self.data_dir,
            "-k", self.socket_dir,
            "-p", str(self.port),
        ]

    def create(self):
        os.makedirs(self.data_dir, exist_ok=True)
        rtn = subprocess.run(
            [
                os.path.join(self.postgres_bin_dir, "initdb"),
                "-D", self.data_dir, #"--auth-local=trust", #"--auth-host"
            ]
        )
        if rtn.returncode != 0:
            raise Exception("initdb failed.")

    def destroy(self):
        self.working_dir.cleanup()

    def up(self, faketime=None):
        self.proc = subprocess.Popen(self.args)

    def down(self):
        if self.proc is None:
            return
        try:
            self.proc.terminate()
        except:  # noqa
            pass

    def initialize(self):
        rtn = subprocess.run(
            [
                os.path.join(self.postgres_bin_dir, "createdb"),
                "-h", self.socket_dir,
                "-p", str(self.port),
                self.db
            ]
        )
        if rtn.returncode != 0:
            raise Exception("createdb failed.")

        rtn = subprocess.run(
            [
                os.path.join(self.postgres_bin_dir, "psql"),
                "-h", self.socket_dir,
                "-p", str(self.port),
                "-d", self.db,
                "-c",
                f"CREATE ROLE {self.user} WITH PASSWORD '{self.password}' SUPERUSER CREATEDB CREATEROLE INHERIT LOGIN;"
            ],
        )
        if rtn.returncode != 0:
            raise Exception("createuser failed.")

    def is_healthy(self):
        rtn = subprocess.run(
            [
                os.path.join(self.postgres_bin_dir, "pg_isready"),
                "-h", self.host,
                "-p", str(self.port),
                "--dbname", self.db,
            ],
            cwd=self.working_dir.name,
        )
        if rtn.returncode != 0:
            return False
        return True

    def wait_for_healthy(self):
        healthy = False
        for trial in range(self.health_check_repeat):
            if self.is_healthy():
                healthy = True
                break
            time.sleep(self.health_check_wait)
        if not healthy:
            raise Exception(f"PostgreSQL service health check failed")

@pytest.fixture(scope="session")
def postgresql_service():
    svc = PostgresqlService()
    svc.create()
    svc.up()
    svc.wait_for_healthy()
    svc.initialize()
    yield svc
    svc.down()
    svc.destroy()

def test_postgresql_query(postgresql_service):
    conn = psycopg2.connect(postgresql_service.dsn)

    try:
        # Prepare data
        cur = conn.cursor()
        cur.execute("create table something(id int, body text)")
        cur.execute("insert into something(id, body) values(1, 'ok')")

        # Execute
        cur.execute("SELECT * FROM something")
        rows = cur.fetchall()

        # Assetion
        assert rows == [(1, 'ok')]
    finally:
        cur.close()
        conn.close()

テストの実装

テストは以下のコマンドで実行できます。

pytest tests/

以上。