A class that wraps a docker image for easier work. Used manily for tests.

container_wrapper.py

import docker
import time
import typing


class ContainerWrapper:
    '''
    A class that wraps a docker image for easier work. Used manily for tests.
    '''
    def __init__(self, image, environment, ports, is_continaer_up: typing.Callable[[], bool]):
        '''
        Creates a new container.
        In order to control it, use the start() and stop() methods.
        '''
        self.container = None
        self.image = image
        self.environment = environment
        self.ports = ports
        self.is_continaer_up = is_continaer_up

    def __enter__(self):
        self.start(timeout=-1)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        self.stop()
        return True

    def __del__(self) -> None:
        self.stop()

    def start(self, timeout: int = 0) -> bool:
        '''
        Starts a new (and empty) container, ready to work on the given port.
        timout - number of seconds to wait. -1 for infinite
        '''
        if self.container:
            raise Exception("Container is already running")
        client = docker.from_env()
        self.container = client.containers.run(self.image, environment=self.environment, detach=True, ports=self.ports)

        # Wait for container to start
        while_timeout = time.time() + timeout
        container_up = False
        while not container_up and (timeout == -1 or time.time() < while_timeout):
            time.sleep(1)
            container_up = self.is_continaer_up()

        return container_up

    def stop(self) -> None:
        '''
        Stops the container.
        '''
        if not self.container:
            return

        self.container.stop()
        self.container.wait()
        self.container = None

    def logs(self, **kwargs) -> typing.Generator[str, None, None]:
        '''
        Get logs from this container. Similar to the ``docker logs`` command.
        See pydocker container for more information
        '''
        if not self.container:
            raise Exception("Container isn't running")

        return self.container.logs(**kwargs)