-
Notifications
You must be signed in to change notification settings - Fork 62
/
sandbox.py
196 lines (155 loc) · 6.36 KB
/
sandbox.py
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
import logging
import os
import shutil
import signal
import socket
from subprocess import Popen, PIPE, CalledProcessError
import sys
import tempfile
import time
import contextlib
import functools
from testify import TestCase, setup, teardown, turtle
from tron.commands import client
# Used for getting the locations of the executable
_test_folder, _ = os.path.split(__file__)
_repo_root, _ = os.path.split(_test_folder)
log = logging.getLogger(__name__)
def wait_on_sandbox(func, delay=0.1, max_wait=5.0):
"""Poll for func() to return True. Sleeps `delay` seconds between polls
up to a max of `max_wait` seconds.
"""
start_time = time.time()
while time.time() - start_time < max_wait:
time.sleep(delay)
if func():
return
raise TronSandboxException("Failed %s" % func.__name__)
def handle_output(cmd, (stdout, stderr), returncode):
"""Log process output before it is parsed. Raise exception if exit code
is nonzero.
"""
if stdout:
log.info("%s: %s", cmd, stdout)
if stderr:
log.warning("%s: %s", cmd, stderr)
if returncode:
raise CalledProcessError(returncode, cmd)
def find_unused_port():
"""Return a port number that is not in use."""
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
with contextlib.closing(sock) as sock:
sock.bind(('localhost', 0))
_, port = sock.getsockname()
return port
class TronSandboxException(Exception):
pass
class SandboxTestCase(TestCase):
_suites = ['sandbox']
@setup
def make_sandbox(self):
self.sandbox = TronSandbox()
@teardown
def delete_sandbox(self):
self.sandbox.delete()
self.sandbox = None
class ClientProxy(object):
"""Wrap calls to client and raise a TronSandboxException on connection
failures.
"""
def __init__(self, client, log_filename):
self.client = client
self.log_filename = log_filename
def log_contents(self):
"""Return the contents of the log file."""
with open(self.log_filename, 'r') as f:
return f.read()
def wrap(self, func, *args, **kwargs):
try:
return func(*args, **kwargs)
except client.RequestError, e:
log.warn("%r, Log:\n%s" % (e, self.log_contents()))
return False
def __getattr__(self, name):
attr = getattr(self.client, name)
if not callable(attr):
return attr
return functools.partial(self.wrap, attr)
class TronSandbox(object):
"""A sandbox for running trond and tron commands in subprocesses."""
def __init__(self):
"""Set up a temp directory and store paths to relevant binaries"""
self.verify_environment()
self.tmp_dir = tempfile.mkdtemp(prefix='tron-')
cmd_path_func = functools.partial(os.path.join, _repo_root, 'bin')
cmds = 'tronctl', 'trond', 'tronfig', 'tronview'
self.commands = dict((cmd, cmd_path_func(cmd)) for cmd in cmds)
self.log_file = self.abs_path('tron.log')
self.log_conf = self.abs_path('logging.conf')
self.pid_file = self.abs_path('tron.pid')
self.config_file = self.abs_path('tron_config.yaml')
self.port = find_unused_port()
self.host = 'localhost'
self.api_uri = 'http://%s:%s' % (self.host, self.port)
client_config = turtle.Turtle(server=self.api_uri,
warn=False, num_displays=100)
cclient = client.Client(client_config)
self.client = ClientProxy(cclient, self.log_file)
self.setup_logging_conf()
def abs_path(self, filename):
"""Return the absolute path for a file in the sandbox."""
return os.path.join(self.tmp_dir, filename)
def setup_logging_conf(self):
config_template = os.path.join(_repo_root, 'tests/data/logging.conf')
with open(config_template, 'r') as fh:
config = fh.read()
with open(self.log_conf, 'w') as fh:
fh.write(config.format(self.log_file))
def verify_environment(self):
ssh_sock = 'SSH_AUTH_SOCK'
msg = "Missing $%s in test environment."
if not os.environ.get(ssh_sock):
raise TronSandboxException(msg % ssh_sock)
path = 'PYTHONPATH'
if not os.environ.get(path):
raise TronSandboxException(msg % path)
def delete(self):
"""Delete the temp directory and shutdown trond."""
if os.path.exists(self.pid_file):
with open(self.pid_file, 'r') as f:
os.kill(int(f.read()), signal.SIGKILL)
shutil.rmtree(self.tmp_dir)
def save_config(self, config_text):
"""Save the initial tron configuration."""
with open(self.config_file, 'w') as f:
f.write(config_text)
def run_command(self, command_name, args=None, stdin_lines=None):
"""Run the command by name and return (stdout, stderr)."""
args = args or []
command = [sys.executable, self.commands[command_name]] + args
stdin = PIPE if stdin_lines else None
proc = Popen(command, stdout=PIPE, stderr=PIPE, stdin=stdin)
streams = proc.communicate(stdin_lines)
handle_output(command, streams, proc.returncode)
return streams
def tronctl(self, args=None):
args = list(args) if args else []
return self.run_command('tronctl', args + ['--server', self.api_uri])
def tronview(self, args=None):
args = list(args) if args else []
args += ['--nocolor', '--server', self.api_uri]
return self.run_command('tronview', args)
def trond(self, args=None):
args = list(args) if args else []
args += ['--working-dir=%s' % self.tmp_dir,
'--pid-file=%s' % self.pid_file,
'--port=%d' % self.port,
'--host=%s' % self.host,
'--config=%s' % self.config_file,
'--log-conf=%s' % self.log_conf]
self.run_command('trond', args)
wait_on_startup = lambda: bool(self.client.home())
wait_on_sandbox(wait_on_startup)
def tronfig(self, config_content):
args = ['--server', self.api_uri, '-']
return self.run_command('tronfig', args, stdin_lines=config_content)