container.py

from __future__ import annotations

import logging
import re
import docker
from time import sleep

LOCALSTACK_IMAGE_NAME = "localstack/localstack-pro"
LATEST_TAG = "latest"

MAX_PORT_CONNECTION_ATTEMPTS = 10
MAX_LOG_COLLECTION_ATTEMPTS = 120
POLL_INTERVAL = 1
NUM_LOG_LINES = 10

ENV_DEBUG = "DEBUG"
ENV_USE_SSL = "USE_SSL"
ENV_DEBUG_DEFAULT = "1"
LOCALSTACK_EXTERNAL_HOSTNAME = "HOSTNAME_EXTERNAL"
DEFAULT_CONTAINER_ID = "localstack-main"

DOCKER_CLIENT = docker.from_env()


class Container:
    @staticmethod
    def create_localstack_container(
        *args,
        pull_new_image: bool = False,
        image_name: str = LOCALSTACK_IMAGE_NAME,
        image_tag: str = LATEST_TAG,
        gateway_listen: str = "0.0.0.0:4566",
        auto_remove: bool = False,
        environment_variables: dict | None = None,
        bind_ports: dict | None = None,
        **kwargs,
    ):
        environment_variables = environment_variables or {}
        environment_variables["GATEWAY_LISTEN"] = gateway_listen

        image_exists = (
            True if len(DOCKER_CLIENT.images.list(name=image_name)) else False
        )

        bind_ports = bind_ports or {}
        gateway_port = gateway_listen.split(":")[1]
        bind_ports.update({gateway_port: gateway_port})

        if pull_new_image or not image_exists:
            logging.info("Pulling latest image")
            DOCKER_CLIENT.images.pull(image_name, image_tag)

        return DOCKER_CLIENT.containers.run(
            image_name,
            ports=bind_ports,
            environment=environment_variables,
            auto_remove=auto_remove,
            detach=True,
        )

    @staticmethod
    def wait_for_ready(container, pattern):
        attempts = 0

        while True:
            logs = container.logs(tail=NUM_LOG_LINES).decode("utf-8")
            if re.search(pattern, logs):
                return

            sleep(POLL_INTERVAL)
            attempts += 1

            if attempts >= MAX_LOG_COLLECTION_ATTEMPTS:
                raise RuntimeError(f"Could not find token: {pattern.pattern} in logs")