1# Original code:
2# Copyright (c) Jupyter Development Team.
3# Distributed under the terms of the Modified BSD License.
4import os
5import logging
6
7import docker
8import pytest
9import requests
10
11from requests.packages.urllib3.util.retry import Retry
12from requests.adapters import HTTPAdapter
13
14
15LOGGER = logging.getLogger(__name__)
16
17@pytest.fixture(scope='session')
18def docker_client():
19    """Docker client configured based on the host environment"""
20    return docker.from_env()
21
22
23@pytest.fixture(scope='session')
24def image_name():
25    """Image name to test"""
26    return os.getenv('TEST_IMAGE', 'afni/afni_dev_base')
27
28
29class TrackedContainer(object):
30    """Wrapper that collects docker container configuration and delays
31    container creation/execution.
32    Parameters
33    ----------
34    docker_client: docker.DockerClient
35        Docker client instance
36    image_name: str
37        Name of the docker image to launch
38    **kwargs: dict, optional
39        Default keyword arguments to pass to docker.DockerClient.containers.run
40    """
41    def __init__(self, docker_client, image_name, **kwargs):
42        self.container = None
43        self.docker_client = docker_client
44        self.image_name = image_name
45        self.kwargs = kwargs
46
47    def run(self, **kwargs):
48        """Runs a docker container using the preconfigured image name
49        and a mix of the preconfigured container options and those passed
50        to this method.
51        Keeps track of the docker.Container instance spawned to kill it
52        later.
53        Parameters
54        ----------
55        **kwargs: dict, optional
56            Keyword arguments to pass to docker.DockerClient.containers.run
57            extending and/or overriding key/value pairs passed to the constructor
58        Returns
59        -------
60        docker.Container
61        """
62        all_kwargs = {}
63        all_kwargs.update(self.kwargs)
64        all_kwargs.update(kwargs)
65        LOGGER.info(f"Running {self.image_name} with args {all_kwargs} ...")
66        self.container = self.docker_client.containers.run(self.image_name, **all_kwargs)
67        return self.container
68
69    def remove(self):
70        """Kills and removes the tracked docker container."""
71        if self.container:
72            self.container.remove(force=True)
73
74
75def get_container(docker_client,image_name):
76    container = TrackedContainer(
77        docker_client,
78        image_name,
79        detach=True,
80    )
81    return container
82
83@pytest.fixture(scope='function')
84def container(docker_client, image_name):
85    """Notebook container with initial configuration appropriate for testing
86    (e.g., HTTP port exposed to the host for HTTP calls).
87    Yields the container instance and kills it when the caller is done with it.
88    """
89    container = get_container(docker_client,image_name)
90    yield container
91    container.remove()
92
93@pytest.fixture(scope='function')
94def named_container(docker_client, request):
95    """
96    Return a container whose image name is defined at the time of test execution
97    """
98    container = get_container(docker_client,request.param)
99    yield container
100    container.remove()
101
102
103