PyTest PostgreSQL Docker image fixture

conftest.py

import time
from typing import Optional

import docker
import pytest


class PostgresImage:
    _container_id: Optional[str] = None
    _container_obj: Optional[docker.models.containers.Container] = None
    _docker_client: docker.client.DockerClient

    def __init__(self):
        self._docker_client = docker.from_env()

    def _wait_until_created(self, tries=10, sleep_seconds=1) -> None:
        for _ in range(tries):
            try:
                self._container_obj = self._docker_client.containers.get(self._container_id)
            except docker.errors.NotFound:
                time.sleep(sleep_seconds)
                continue
            else:
                break
        else:
            raise Exception("Container not created")

    def _check_ready(self, tries=30, sleep_seconds=1) -> None:
        ready = False
        for _ in range(tries):
            # Consider the container ready when 'database system is ready to accept connections' is seen AFTER a
            # 'listening on IPv4 address "0.0.0.0", port 5432' line. Those before that are initialization by the image's
            # entrypoint scripts.
            found_network_activated_line = False
            for line in self._container_obj.logs().decode('utf-8').splitlines():
                if 'listening on IPv4 address' in line:
                    found_network_activated_line = True
                if found_network_activated_line and 'database system is ready to accept connections' in line:
                    ready = True
                    break
            if ready:
                break
            else:
                time.sleep(sleep_seconds)
                continue
        else:
            raise Exception("Container not ready within timeout")

    def run(self):
        container_image = "postgres:13"
        image_options = dict(
            mem_limit='1g',
            environment={
                "POSTGRES_PASSWORD": "pytest",
                "POSTGRES_DB": "pytest",
                "POSTGRES_USER": "pytest",
            },
            privileged=False,
            detach=True,
            tmpfs=["/var/lib/postgresql/data"],  # No persistence required, create tempdir in memory.
            remove=True,  # clean up container after it has run
        )
        container = self._docker_client.containers.run(container_image, **image_options)
        self._container_id = container.id
        self._wait_until_created()
        self._check_ready()
        return self._container_obj.attrs['NetworkSettings']['IPAddress']

    def stop(self):
        if self._container_id is None:
            return
        try:
            self._wait_until_created()  # In case it's stopped while it was still creating.
        except:
            return
        try:
            self._container_obj.kill()
        except docker.errors.APIError:
            pass


@pytest.fixture(scope='session')
def postgres_container() -> str:
    postgres_image = PostgresImage()
    host = postgres_image.run()
    yield "postgresql+asyncpg://pytest:pytest@%s/pytest" % host
    postgres_image.stop()