aboutsummaryrefslogtreecommitdiffhomepage
path: root/src/test/docker/docker_test.py
blob: a185e277755ff336ec5a8d508e1180d5ec456f0b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
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)))