Skip to content

Instantly share code, notes, and snippets.

@gertvdijk
Created March 15, 2021 13:39
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save gertvdijk/ee9dd9b980197dd660282c030ab512f1 to your computer and use it in GitHub Desktop.
Save gertvdijk/ee9dd9b980197dd660282c030ab512f1 to your computer and use it in GitHub Desktop.
PyTest PostgreSQL Docker image fixture
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment