| #!/usr/bin/env python3 |
| # -*- coding: utf-8 -*- |
| |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """ |
| Safe docker run tests |
| """ |
| import itertools |
| import os |
| import signal |
| import unittest |
| from typing import Optional |
| from unittest.mock import create_autospec, patch, call |
| |
| from docker import DockerClient |
| from docker.models.containers import Container, ContainerCollection |
| |
| from safe_docker_run import SafeDockerClient, main |
| |
| |
| def create_mock_container(status_code: int = 0): |
| """ |
| Creates a mock docker container that exits with the specified status code |
| """ |
| mock_container = create_autospec(Container, name="mock_container") |
| mock_container.wait.return_value = { |
| "StatusCode": status_code |
| } |
| return mock_container |
| |
| |
| def create_mock_container_collection(container: Container): |
| """ |
| Creates a mock ContainerCollection that return the supplied container when the 'run' method is called |
| """ |
| mock_container_collection = create_autospec(ContainerCollection, name="mock_collection") |
| mock_container_collection.run.return_value = container |
| return mock_container_collection |
| |
| |
| class MockDockerClient: |
| """ |
| A mock DockerClient when docker.from_env is called |
| The supplied container will be returned when the client.containers.run method is called |
| """ |
| def __init__(self, container: Container): |
| self._mock_client = create_autospec(DockerClient, name="mock_client") |
| self._mock_client.containers = create_mock_container_collection(container) |
| self._patch = patch("docker.from_env", return_value=self._mock_client) |
| |
| def __enter__(self): |
| self._patch.start() |
| return self._mock_client |
| |
| def __exit__(self, _, __, ___): |
| self._patch.stop() |
| |
| |
| class TestSafeDockerRun(unittest.TestCase): |
| |
| @patch("safe_docker_run.signal.pthread_sigmask") |
| @patch.dict(os.environ, { |
| "BUILD_NUMBER": "BUILD_NUMBER_5", |
| "BUILD_ID": "BUILD_ID_1", |
| "BUILD_TAG": "BUILD_TAG_7" |
| }) |
| def test_run_successful(self, mock_pthread_sigmask): |
| """ |
| Tests successful run |
| """ |
| mock_container = create_mock_container() |
| |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| |
| # Check return code is 0 |
| assert safe_docker.run("image", "command") == 0 |
| |
| # Check call to container is correct |
| assert mock_client.containers.run.call_args_list == [ |
| call("image", "command", detach=True, environment={ |
| "BUILD_NUMBER": "BUILD_NUMBER_5", |
| "BUILD_ID": "BUILD_ID_1", |
| "BUILD_TAG": "BUILD_TAG_7" |
| }) |
| ] |
| |
| # Check correct signals are blocked then unblocked |
| assert mock_pthread_sigmask.call_args_list == [ |
| call(signal.SIG_BLOCK, {signal.SIGINT, signal.SIGTERM}), |
| call(signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM}) |
| ] |
| |
| # Assert container is stopped and removed |
| assert mock_container.stop.call_count == 1 |
| assert mock_container.remove.call_count == 1 |
| assert len(safe_docker._containers) == 0 |
| |
| def test_run_detach(self): |
| """ |
| Tests detach=True is passed to the underlying call by default |
| """ |
| mock_container = create_mock_container() |
| |
| # Test detach=True is passed in even if not specified |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 0 |
| assert mock_client.containers.run.call_count == 1 |
| _, kwargs = mock_client.containers.run.call_args |
| assert kwargs["detach"] is True |
| |
| # Test passing in detach=True does not cause any issues |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command", detach=True) == 0 |
| assert mock_client.containers.run.call_count == 1 |
| _, kwargs = mock_client.containers.run.call_args |
| assert kwargs["detach"] is True |
| |
| # Test detach=False fails |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| with self.assertRaises(ValueError): |
| safe_docker.run("image", "command", detach=False) |
| assert mock_client.containers.run.call_args_list == [] |
| |
| def test_jenkins_vars(self): |
| """ |
| Tests jenkins environment variables are appropriately passed to the underlying docker run call |
| """ |
| # NOTE: It's important that these variables are passed to the underlying docker container |
| # These variables are passed to the container so the process tree killer can find runaway |
| # process inside the container |
| # https://wiki.jenkins.io/display/JENKINS/ProcessTreeKiller |
| # https://github.com/jenkinsci/jenkins/blob/578d6bacb33a5e99f149de504c80275796f0b231/core/src/main/java/hudson/model/Run.java#L2393 |
| |
| jenkins_vars = { |
| "BUILD_NUMBER": "BUILD_NUMBER_5", |
| "BUILD_ID": "BUILD_ID_1", |
| "BUILD_TAG": "BUILD_TAG_7" |
| } |
| mock_container = create_mock_container() |
| |
| # Test environment is empty if the jenkins vars are not present |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 0 |
| assert mock_client.containers.run.call_count == 1 |
| _, kwargs = mock_client.containers.run.call_args |
| assert kwargs["environment"] == {} |
| |
| # Test environment contains jenkins env vars if they are present |
| with MockDockerClient(mock_container) as mock_client: |
| with patch.dict(os.environ, jenkins_vars): |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 0 |
| assert mock_client.containers.run.call_count == 1 |
| _, kwargs = mock_client.containers.run.call_args |
| assert kwargs["environment"] == jenkins_vars |
| |
| # Test jenkins env vars are added to callers env vars |
| user_env = {"key1": "value1", "key2": "value2"} |
| with MockDockerClient(mock_container) as mock_client: |
| with patch.dict(os.environ, jenkins_vars): |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command", environment=user_env) == 0 |
| assert mock_client.containers.run.call_count == 1 |
| _, kwargs = mock_client.containers.run.call_args |
| assert kwargs["environment"] == {**jenkins_vars, **user_env} |
| |
| def test_run_args_kwargs_passed(self): |
| """ |
| Tests args and kwargs are passed to the container run call |
| """ |
| mock_container = create_mock_container() |
| |
| # Test detach=True is passed in even if not specified |
| with MockDockerClient(mock_container) as mock_client: |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run( |
| "image", |
| "command", |
| "another_arg", |
| str_param="value", |
| bool_param=True, |
| none_param=None, |
| int_param=5, |
| float_param=5.2, |
| list_param=["this", "is", "a", "list"], |
| map_param={ |
| "a": "5", |
| "b": True, |
| "c": 2 |
| }) == 0 |
| assert mock_client.containers.run.call_args_list == [ |
| call( |
| "image", |
| "command", |
| "another_arg", |
| detach=True, |
| environment={}, |
| str_param="value", |
| bool_param=True, |
| none_param=None, |
| int_param=5, |
| float_param=5.2, |
| list_param=["this", "is", "a", "list"], |
| map_param={ |
| "a": "5", |
| "b": True, |
| "c": 2 |
| } |
| ) |
| ] |
| |
| def test_container_returns_non_zero_status_code(self): |
| """ |
| Tests non-zero code from container is returned and the container |
| is cleaned up |
| """ |
| mock_container = create_mock_container(status_code=10) |
| with MockDockerClient(mock_container): |
| safe_docker = SafeDockerClient() |
| # check return code and that container gets cleaned up |
| assert safe_docker.run("image", "command") == 10 |
| assert mock_container.stop.call_count == 1 |
| assert mock_container.remove.call_count == 1 |
| assert len(safe_docker._containers) == 0 |
| |
| def test_container_wait_raises_returns_150(self): |
| """ |
| Tests 150 is returned if an error is raised when calling container.wait |
| """ |
| mock_container = create_mock_container() |
| mock_container.wait.side_effect = RuntimeError("Something bad happened") |
| with MockDockerClient(mock_container): |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 150 |
| |
| def test_container_stop_raises_returns_151(self): |
| """ |
| Tests 151 is returned if an error is raised when calling container.stop |
| """ |
| mock_container = create_mock_container() |
| mock_container.stop.side_effect = RuntimeError("Something bad happened") |
| with MockDockerClient(mock_container): |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 151 |
| |
| def test_container_remove_raises_returns_152(self): |
| """ |
| Tests 152 is returned if an error is raised when calling container.remove |
| """ |
| mock_container = create_mock_container() |
| mock_container.remove.side_effect = RuntimeError("Something bad happened") |
| with MockDockerClient(mock_container): |
| safe_docker = SafeDockerClient() |
| assert safe_docker.run("image", "command") == 152 |
| |
| def test_main(self): |
| """ |
| Tests main function against different command line arguments |
| """ |
| tests = [ |
| # ( supplied command line arguments, expected call ) |
| ( |
| ["image", "command"], |
| call("image", command="command", runtime=None, user=None, name=None, volumes={}, cap_add=[]) |
| ), |
| ( |
| ["image", "command", "arg1", "arg2"], |
| call("image", command="command arg1 arg2", runtime=None, user=None, name=None, volumes={}, cap_add=[]) |
| ), |
| ( |
| ["--runtime", "nvidia", "image", "command"], |
| call("image", command="command", runtime="nvidia", user=None, name=None, volumes={}, cap_add=[]) |
| ), |
| ( |
| ["--user", "1001:1001", "image", "command"], |
| call("image", command="command", runtime=None, user="1001:1001", name=None, volumes={}, cap_add=[]) |
| ), |
| ([ |
| "--volume", "/local/path1:/container/path1", |
| "--volume", "/local/path2:/container/path2:ro", |
| "image", |
| "command" |
| ], call("image", command="command", runtime=None, user=None, name=None, volumes={ |
| "/local/path1": { |
| "bind": "/container/path1", |
| "mode": "rw" |
| }, |
| "/local/path2": { |
| "bind": "/container/path2", |
| "mode": "ro" |
| } |
| }, cap_add=[])), |
| ([ |
| "--runtime", "nvidia", |
| "-u", "1001:1001", |
| "-v", "/local/path1:/container/path1", |
| "-v", "/local/path2:/container/path2:ro", |
| "--cap-add", "bob", |
| "--cap-add", "jimmy", |
| "--name", |
| "container_name", |
| "image", |
| "command", |
| "arg1", |
| "arg2" |
| ], call( |
| "image", |
| command="command arg1 arg2", |
| runtime="nvidia", |
| user="1001:1001", |
| name="container_name", |
| volumes={ |
| "/local/path1": { |
| "bind": "/container/path1", |
| "mode": "rw" |
| }, |
| "/local/path2": { |
| "bind": "/container/path2", |
| "mode": "ro" |
| } |
| }, cap_add=["bob", "jimmy"]) |
| ) |
| ] |
| |
| # Tests valid arguments |
| mock_docker = create_autospec(SafeDockerClient) |
| mock_docker.run.return_value = 0 |
| with patch("safe_docker_run.SafeDockerClient", return_value=mock_docker): |
| for test in tests: |
| arguments, expected_call = test |
| main(arguments) |
| assert mock_docker.run.call_args == expected_call |
| |
| # Tests invalid arguments |
| tests = [ |
| [], |
| None, |
| ["image"], |
| # Test some bad volume mounts |
| ["-v", "bob", "image", "args"], |
| ["-v", "/local/path", "image", "args"], |
| ["-v", "/local/path:/container/path:blah", "image", "args"], |
| ["-v", "", "image", "args"], |
| ["-v", "a:b:c:d", "image", "args"] |
| ] |
| |
| mock_docker = create_autospec(SafeDockerClient) |
| with patch("safe_docker_run.SafeDockerClient", return_value=mock_docker): |
| with self.assertRaises(SystemExit): |
| for test in tests: |
| main(test) |
| |
| def test_clean_up(self): |
| """ |
| Tests container clean up in case of SIGTERM and SIGINT |
| """ |
| import subprocess |
| import time |
| import docker.errors |
| |
| docker_client = docker.from_env() |
| container_name = "safedockertestcontainer1234" |
| |
| def get_container(name: str) -> Optional[Container]: |
| try: |
| return docker_client.containers.get(name) |
| except docker.errors.NotFound: |
| return None |
| |
| def remove_container_if_exists(name: str): |
| container = get_container(name) |
| if container: |
| container.stop() |
| container.remove() |
| |
| def wait_for_container(name: str) -> bool: |
| for _ in itertools.count(5): |
| if get_container(name): |
| return True |
| time.sleep(1) |
| return False |
| |
| # Clear any containers with container name |
| remove_container_if_exists(container_name) |
| |
| # None => not signal is emitted - we should still finish with no containers at the end due |
| # to the atexit |
| for sig in [None, signal.SIGTERM, signal.SIGINT]: |
| # Execute the safe docker run script in a different process |
| proc = subprocess.Popen(['python3','safe_docker_run.py', "--name", container_name, "ubuntu:18.04", "sleep 10"]) |
| # NOTE: we need to wait for the container to come up as not all operating systems support blocking signals |
| if wait_for_container(container_name) is False: |
| raise RuntimeError("Test container did not come up") |
| |
| # Issue the signal and wait for the process to finish |
| if sig: |
| proc.send_signal(sig) |
| proc.wait() |
| |
| # The container should no longer exist |
| assert get_container(container_name) is None |
