| # Copyright 2016 The Bazel Authors. All rights reserved. |
| # |
| # Licensed 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. |
| """A simple test runner for docker (experimental).""" |
| |
| import copy |
| import os |
| import os.path |
| import shlex |
| import StringIO |
| import subprocess |
| import sys |
| import threading |
| |
| from third_party.py import gflags |
| |
| gflags.DEFINE_multistring( |
| "image", [], |
| "The list of additional docker image to load (path to a docker_build " |
| "target), optional.") |
| |
| gflags.DEFINE_string( |
| "main", None, |
| "The main image to run (path to a docker_build target), mandatory.") |
| gflags.MarkFlagAsRequired("main") |
| |
| gflags.DEFINE_string( |
| "cmd", None, |
| "A command to provide to the docker image, optional (default: use the " |
| "entrypoint).") |
| |
| gflags.DEFINE_string("docker", "docker", "Path to the docker client binary.") |
| |
| gflags.DEFINE_boolean("verbose", True, "Be verbose.") |
| |
| FLAGS = gflags.FLAGS |
| |
| LOCAL_IMAGE_PREFIX = "bazel/" |
| |
| |
| def _copy_stream(in_stream, out_stream): |
| for c in iter(lambda: in_stream.read(1), ""): |
| out_stream.write(c) |
| out_stream.flush() |
| |
| |
| def execute(command, stdout=sys.stdout, stderr=sys.stderr, env=None): |
| """Execute a command while redirecting its output streams.""" |
| if FLAGS.verbose: |
| print "Executing '%s'" % " ".join(command) |
| p = subprocess.Popen(command, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| env=env) |
| t1 = threading.Thread(target=_copy_stream, args=[p.stdout, stdout]) |
| t2 = threading.Thread(target=_copy_stream, args=[p.stderr, stderr]) |
| t1.daemon = True |
| t2.daemon = True |
| t1.start() |
| t2.start() |
| p.wait() |
| t1.join() |
| t2.join() |
| return p.returncode |
| |
| |
| def load_image(image): |
| """Load a docker image using the runner provided by docker_build.""" |
| tag = LOCAL_IMAGE_PREFIX + ":".join(image.rsplit("/", 1)) |
| err = StringIO.StringIO() |
| env = copy.deepcopy(os.environ) |
| env["DOCKER"] = FLAGS.docker |
| ret = execute([image], stderr=err, env=env) |
| if ret != 0: |
| sys.stderr.write("Error loading image %s (return code: %s):\n" % |
| (image, ret)) |
| sys.stderr.write(err.getvalue()) |
| return None |
| return tag |
| |
| |
| def load_images(images): |
| """Load a series of docker images using the docker_build's runner.""" |
| print "### Image loading ###" |
| return [load_image(image) for image in images] |
| |
| |
| def cleanup_images(tags): |
| """Remove docker tags and images previously loaded.""" |
| print "### Image cleanup ###" |
| for tag in tags: |
| if isinstance(tag, basestring): |
| execute([FLAGS.docker, "rmi", tag]) |
| |
| |
| def run_image(tag): |
| """Run a docker image, in background.""" |
| print "Running " + tag |
| out = StringIO.StringIO() |
| err = StringIO.StringIO() |
| process = execute([FLAGS.docker, "run", "--rm", tag], out, err) |
| if process.wait() != 0: |
| sys.stderr.write("Error running docker run on %s:\n" % tag) |
| sys.stderr.write(err.getvalue()) |
| return None |
| else: |
| return out.getvalue().strip() |
| |
| |
| def run_images(tags): |
| """Run a list of docker images, in background.""" |
| print "### Running images ###" |
| return [run_image(tag) for tag in tags] |
| |
| |
| def cleanup_containers(containers): |
| """Kill containers.""" |
| print "### Containers cleanup ###" |
| for c in containers: |
| if isinstance(c, basestring): |
| execute([FLAGS.docker, "kill", c]) |
| |
| |
| def main(unused_argv): |
| tags = load_images([FLAGS.main] + FLAGS.image) |
| if None in tags: |
| cleanup_images(tags) |
| return -1 |
| try: |
| containers = run_images(tags[1:]) |
| ret = -1 |
| if None not in containers: |
| print "### Running main container ###" |
| ret = execute([ |
| FLAGS.docker, |
| "run", |
| "--rm", |
| "--attach=STDOUT", |
| "--attach=STDERR", tags[0] |
| ] + ([] if FLAGS.cmd is None else shlex.split(FLAGS.cmd))) |
| finally: |
| cleanup_containers(containers) |
| cleanup_images(tags) |
| return ret |
| |
| |
| if __name__ == "__main__": |
| sys.exit(main(FLAGS(sys.argv))) |