1 files changed, 158 insertions, 0 deletions
diff --git a/src/test/docker/docker_test.py b/src/test/docker/docker_test.py
new file mode 100644
index 0000000000..a185e27775
--- /dev/null
+++ b/src/test/docker/docker_test.py
@@ -0,0 +1,158 @@
+# 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)))