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