Merge lp:~sseman/juju-chaos-monkey/upstart-init into lp:juju-chaos-monkey

Proposed by Seman
Status: Merged
Merged at revision: 19
Proposed branch: lp:~sseman/juju-chaos-monkey/upstart-init
Merge into: lp:juju-chaos-monkey
Diff against target: 896 lines (+380/-207)
14 files modified
chaos/kill.py (+8/-3)
chaos/net.py (+0/-4)
chaos_monkey.py (+0/-30)
chaos_monkey_base.py (+0/-4)
runner.py (+56/-29)
scripts/chaos-monkey-restart.conf (+9/-0)
scripts/restart_chaos_monkey.py (+28/-0)
tests/test_chaos_monkey.py (+0/-84)
tests/test_init.py (+87/-0)
tests/test_kill.py (+1/-1)
tests/test_net.py (+0/-6)
tests/test_restart_chaos_monkey.py (+35/-0)
tests/test_runner.py (+82/-46)
utils/init.py (+74/-0)
To merge this branch: bzr merge lp:~sseman/juju-chaos-monkey/upstart-init
Reviewer Review Type Date Requested Status
John George (community) Approve
Review via email: mp+260967@code.launchpad.net

Description of the change

This branch adds a support for restarting chaos monkey after a machine reboot. It writes init script to /etc/init before restarting a unit and removes it after a unit has rebooted.

To post a comment you must log in.
20. By Seman

Updated the init script and added a unit test.

Revision history for this message
Seman (sseman) wrote :

Adding sample of init script generated by the Chaos Monkey:

------
data: description "Restarts Chaos Monkey."
start on runlevel [2345]
stop on runlevel [!2345]
normal exit 0

task
script
  exec python /home/ubuntu/cm/scripts/restart_chaos_monkey.py --runner-path '/home/ubuntu/cm/runner.py' --cmd-arg '--include-command restart-unit --total-timeout 10 --enablement-timeout 1 /home/ubuntu' --expire-time 1433401733.06
end script
-------

Revision history for this message
Seman (sseman) wrote :

During a reboot, the upstart init script executes a Python code in Chaos Monkey "scripts/restart_chaos_monkey.py". The code, restart_chaos_monkey.py, is responsible for restarting and continuing executing of Chaos Monkey with the same command line arguments the user initially started with.

Revision history for this message
John George (jog) :
21. By Seman

Removed os.system and shutdown method.

Revision history for this message
Seman (sseman) wrote :

Removed os.system() and shutdown() method. The shutdown() method is not necessary since the code has been updated to make each net command self-contained and calling the shutdown('ufw --force reset') after executing system level reboot request is causing an issue.

Revision history for this message
John George (jog) :
Revision history for this message
John George (jog) wrote :

comment in line.

22. By Seman

Added expire-time to the random_chaos() call.

Revision history for this message
John George (jog) :
review: Needs Fixing
23. By Seman

Updated the flag for the lock file.

Revision history for this message
Seman (sseman) wrote :

Made updates. Please see in-line comments.

Revision history for this message
John George (jog) wrote :

Thank you for the updates.

review: Approve

Preview Diff

[H/L] Next/Prev Comment, [J/K] Next/Prev File, [N/P] Next/Prev Hunk
1=== modified file 'chaos/kill.py'
2--- chaos/kill.py 2015-05-21 23:55:13 +0000
3+++ chaos/kill.py 2015-06-05 00:54:01 +0000
4@@ -18,6 +18,7 @@
5
6 jujud_cmd = 'kill-jujud'
7 mongod_cmd = 'kill-mongod'
8+ restart_cmd = 'restart-unit'
9 group = 'kill'
10
11 def __init__(self):
12@@ -75,7 +76,11 @@
13 group=self.group,
14 command_str=self.mongod_cmd,
15 description='Kill mongod process.'))
16+ chaos.append(
17+ Chaos(
18+ enable=self.restart_unit,
19+ disable=None,
20+ group=self.group,
21+ command_str=self.restart_cmd,
22+ description='Restart the unit.'))
23 return chaos
24-
25- def shutdown(self):
26- pass
27
28=== modified file 'chaos/net.py'
29--- chaos/net.py 2015-05-28 15:35:16 +0000
30+++ chaos/net.py 2015-06-05 00:54:01 +0000
31@@ -119,7 +119,3 @@
32 allow_in_to_any,
33 ),
34 ]
35-
36- def shutdown(self):
37- # TODO(gz): Delete this? Each Chaos object cleans up firewall rules.
38- run_shell_command("ufw --force reset")
39
40=== modified file 'chaos_monkey.py'
41--- chaos_monkey.py 2015-05-16 02:25:03 +0000
42+++ chaos_monkey.py 2015-06-05 00:54:01 +0000
43@@ -1,12 +1,7 @@
44-import logging
45-import random
46-from time import sleep
47-
48 from chaos import (
49 kill,
50 net
51 )
52-from utility import NotFound
53
54 __metaclass__ = type
55
56@@ -90,28 +85,3 @@
57 if item.command_str == command_str:
58 return item
59 return None
60-
61- def run_random_chaos(self, timeout=2):
62- random_chaos = random.choice(self.chaos)
63- self._run_command(random_chaos, timeout=timeout)
64-
65- def run_chaos(self, group, command_str, timeout=2):
66- command_found = False
67- for chaos in self.chaos:
68- if chaos.group == group and chaos.command_str == command_str:
69- command_found = True
70- self._run_command(chaos, timeout=timeout)
71- if not command_found:
72- raise NotFound("Command not found: group: %s command_str:%s" % (
73- group, command_str))
74-
75- def _run_command(self, chaos, timeout=2):
76- logging.info("{}".format(chaos.description))
77- chaos.enable()
78- sleep(timeout)
79- if chaos.disable:
80- chaos.disable()
81-
82- def shutdown(self):
83- for obj in self.factory_obj:
84- obj.shutdown()
85
86=== modified file 'chaos_monkey_base.py'
87--- chaos_monkey_base.py 2015-05-08 04:53:12 +0000
88+++ chaos_monkey_base.py 2015-06-05 00:54:01 +0000
89@@ -13,10 +13,6 @@
90 def get_chaos(self):
91 raise NotImplemented
92
93- @abc.abstractmethod
94- def shutdown(self):
95- raise NotImplemented
96-
97
98 class Chaos:
99
100
101=== modified file 'runner.py'
102--- runner.py 2015-05-22 19:53:05 +0000
103+++ runner.py 2015-06-05 00:54:01 +0000
104@@ -2,10 +2,15 @@
105 import errno
106 import logging
107 import os
108+import random
109 import signal
110 import sys
111-from time import time
112+from time import (
113+ time,
114+ sleep
115+)
116
117+from chaos.kill import Kill
118 from chaos_monkey import ChaosMonkey
119 from utility import (
120 BadRequest,
121@@ -14,6 +19,7 @@
122 setup_logging,
123 split_arg_string,
124 )
125+from utils.init import Init
126
127
128 class Runner:
129@@ -25,6 +31,7 @@
130 self.workspace_lock = False
131 self.lock_file = '{}/{}'.format(self.workspace, 'chaos_runner.lock')
132 self.chaos_monkey = chaos_monkey
133+ self.expire_time = None
134
135 @classmethod
136 def factory(cls, workspace, log_count=1, dry_run=False):
137@@ -35,13 +42,16 @@
138 chaos_monkey = ChaosMonkey.factory()
139 return cls(workspace, chaos_monkey, log_count, dry_run)
140
141- def acquire_lock(self):
142+ def acquire_lock(self, restart=False):
143 if not os.path.isdir(self.workspace):
144 sys.stderr.write('Not a directory: {}\n'.format(self.workspace))
145 sys.exit(-1)
146+ init = Init.upstart()
147+ init.uninstall()
148 try:
149- lock_fd = os.open(self.lock_file,
150- os.O_CREAT | os.O_EXCL | os.O_WRONLY)
151+ file_flag = ((os.O_CREAT | os.O_EXCL | os.O_WRONLY)
152+ if not restart else (os.O_CREAT | os.O_WRONLY))
153+ lock_fd = os.open(self.lock_file, file_flag)
154 except OSError as e:
155 if e.errno == errno.EEXIST:
156 sys.stderr.write('Lock file already exists: {}\n'.format(
157@@ -65,26 +75,34 @@
158
159 def random_chaos(self, run_timeout, enablement_timeout, include_group=None,
160 exclude_group=None, include_command=None,
161- exclude_command=None, run_once=False):
162- """Runs a random chaos monkey.
163-
164- :param run_timeout: Total time to run the chaos
165- :param enablement_timeout: Time between enabling and disabling chaos.
166- Example: disable all the network, wait for timeout and enable it back
167- :rtype none
168- """
169+ exclude_command=None, run_once=False, expire_time=None):
170+ """Runs a random chaos monkey."""
171 self.filter_commands(
172 include_group=include_group, exclude_group=exclude_group,
173 include_command=include_command, exclude_command=exclude_command)
174- expire_time = time() + run_timeout
175- while time() < expire_time:
176+ self.expire_time = expire_time or (time() + run_timeout)
177+ while time() < self.expire_time:
178 if self.stop_chaos or self.dry_run:
179 break
180- self.chaos_monkey.run_random_chaos(enablement_timeout)
181+ self._run_command(enablement_timeout)
182 if run_once:
183 break
184- if not self.dry_run:
185- self.chaos_monkey.shutdown()
186+
187+ def _run_command(self, enablement_timeout):
188+ chaos = random.choice(self.chaos_monkey.chaos)
189+ logging.info("{}".format(chaos.description))
190+ if chaos.command_str == Kill.restart_cmd:
191+ self.stop_chaos = True
192+ init = Init.upstart()
193+ init.install(
194+ cmd_arg=' '.join(sys.argv[1:]), expire_time=self.expire_time)
195+ chaos.enable()
196+ if chaos.command_str == Kill.restart_cmd:
197+ return
198+
199+ sleep(enablement_timeout)
200+ if chaos.disable:
201+ chaos.disable()
202
203 def cleanup(self):
204 if self.lock_file:
205@@ -95,7 +113,7 @@
206 raise
207 logging.warning('Lock file not found: {}'.format(
208 self.lock_file))
209- logging.info('Chaos monkey stopped')
210+ logging.info('Chaos Monkey stopped.\n')
211
212 def filter_commands(self, include_group=None, exclude_group=None,
213 include_command=None, exclude_command=None):
214@@ -209,19 +227,26 @@
215 parser.add_argument(
216 '-ro', '--run-once', action='store_true',
217 help='Run a single command only.', default=False)
218+ parser.add_argument(
219+ '-r', '--restart', action='store_true',
220+ help='Indicates the run is after a reboot.', default=False)
221+ parser.add_argument(
222+ '-ep', '--expire-time', type=float,
223+ help='Chaos Monkey expire time (UNIX timestamp).', default=None)
224 args = parser.parse_args(argv)
225
226 if args.run_once and args.total_timeout:
227 parser.error("Conflicting request: total-timeout is irrelevant "
228 "if run-once is set.")
229- if not args.total_timeout:
230- args.total_timeout = args.enablement_timeout
231- if args.enablement_timeout > args.total_timeout:
232- parser.error("total-timeout can not be less than "
233- "enablement-timeout.")
234- if args.total_timeout <= 0:
235- parser.error("Invalid total-timeout value: timeout must be "
236- "greater than zero.")
237+ if not args.expire_time:
238+ if not args.total_timeout:
239+ args.total_timeout = args.enablement_timeout
240+ if args.enablement_timeout > args.total_timeout:
241+ parser.error("total-timeout can not be less than "
242+ "enablement-timeout.")
243+ if args.total_timeout <= 0:
244+ parser.error("Invalid total-timeout value: timeout must be "
245+ "greater than zero.")
246 if args.enablement_timeout < 0:
247 parser.error("Invalid enablement-timeout value: timeout must be "
248 "zero or greater.")
249@@ -232,9 +257,10 @@
250 runner = Runner.factory(workspace=args.path, log_count=args.log_count,
251 dry_run=args.dry_run)
252 setup_sig_handlers(runner.sig_handler)
253- runner.acquire_lock()
254- logging.info('Chaos monkey started in {}'.format(args.path))
255+ msg = 'started' if not args.restart else 'restarted after a reboot'
256+ logging.info('Chaos Monkey {} in {}'.format(msg, args.path))
257 logging.debug('Dry run is set to {}'.format(args.dry_run))
258+ runner.acquire_lock(restart=args.restart)
259 try:
260 runner.random_chaos(
261 run_timeout=args.total_timeout,
262@@ -243,7 +269,8 @@
263 exclude_group=args.exclude_group,
264 include_command=args.include_command,
265 exclude_command=args.exclude_command,
266- run_once=args.run_once)
267+ run_once=args.run_once,
268+ expire_time=args.expire_time)
269 except Exception as e:
270 logging.error('{} ({})'.format(e, type(e).__name__))
271 sys.exit(1)
272
273=== added directory 'scripts'
274=== added file 'scripts/__init__.py'
275=== added file 'scripts/chaos-monkey-restart.conf'
276--- scripts/chaos-monkey-restart.conf 1970-01-01 00:00:00 +0000
277+++ scripts/chaos-monkey-restart.conf 2015-06-05 00:54:01 +0000
278@@ -0,0 +1,9 @@
279+description "Restarts Chaos Monkey."
280+start on runlevel [2345]
281+stop on runlevel [!2345]
282+normal exit 0
283+
284+task
285+script
286+ exec python {restart_script_path} --runner-path '{runner_path}' --cmd-arg '{cmd_arg}' --expire-time {expire_time}
287+end script
288
289=== added file 'scripts/restart_chaos_monkey.py'
290--- scripts/restart_chaos_monkey.py 1970-01-01 00:00:00 +0000
291+++ scripts/restart_chaos_monkey.py 2015-06-05 00:54:01 +0000
292@@ -0,0 +1,28 @@
293+from argparse import ArgumentParser
294+import subprocess
295+
296+
297+def parse_args(argv=None):
298+ parser = ArgumentParser()
299+ parser.add_argument(
300+ '--runner-path', help='Chaos Monkey runner path.', default=None)
301+ parser.add_argument(
302+ '--expire-time', help='Chaos Monkey expire time.', default=None,
303+ type=float)
304+ parser.add_argument(
305+ '--cmd-arg', help='Chaos Monkey command arguments.', default=None)
306+ args = parser.parse_args(argv)
307+ if not args.runner_path or not args.expire_time or not args.cmd_arg:
308+ parser.error("Invalid command arguments.")
309+ return args
310+
311+
312+def restart_chaos_monkey(args):
313+ cmd = (['python'] + [args.runner_path] + args.cmd_arg.split(' ') +
314+ ['--expire-time'] + [str(args.expire_time)] + ['--restart'])
315+ subprocess.Popen(cmd)
316+
317+
318+if __name__ == '__main__':
319+ args = parse_args()
320+ restart_chaos_monkey(args)
321
322=== modified file 'tests/test_chaos_monkey.py'
323--- tests/test_chaos_monkey.py 2015-05-29 11:33:58 +0000
324+++ tests/test_chaos_monkey.py 2015-06-05 00:54:01 +0000
325@@ -1,11 +1,7 @@
326-from mock import patch, call
327-
328 from chaos_monkey import (
329 ChaosMonkey,
330- NotFound
331 )
332 from chaos.kill import Kill
333-from chaos.net import Net
334 from tests.common_test_base import CommonTestBase
335 from tests.test_kill import get_all_kill_commands
336 from tests.test_net import get_all_net_commands
337@@ -28,86 +24,6 @@
338 self.assertItemsEqual(
339 self._get_command_str(all_chaos), self._get_all_command_strings())
340
341- def test_run_random_chaos(self):
342- cm = ChaosMonkey.factory()
343- cm.include_group('all')
344- with patch('utility.check_output', autospec=True) as mock:
345- cm.run_random_chaos(timeout=0)
346- self.assertEqual(mock.called, True)
347-
348- def test_run_random_chaos_passes_timeout(self):
349- cm = ChaosMonkey.factory()
350- cm.include_group('all')
351- with patch('chaos_monkey.ChaosMonkey._run_command',
352- autospec=True) as mock:
353- cm.run_random_chaos(timeout=1)
354- self.assertEqual(1, mock.call_args_list[0][1]['timeout'])
355-
356- def test_run_chaos(self):
357- cm = ChaosMonkey.factory()
358- cm.include_group('all')
359- with patch('utility.check_output', autospec=True) as mock:
360- cm.run_chaos('net', 'deny-state-server', timeout=0)
361- self.assertEqual(mock.mock_calls, [
362- call(['ufw', 'deny', '37017']),
363- call(['ufw', 'allow', 'in', 'to', 'any']),
364- call(['ufw', '--force', 'enable']),
365- call(['ufw', 'disable']),
366- call(['ufw', 'delete', 'allow', 'in', 'to', 'any']),
367- call(['ufw', 'delete', 'deny', '37017']),
368- ])
369-
370- def test_run_chaos_passes_timeout(self):
371- cm = ChaosMonkey.factory()
372- cm.include_group('all')
373- with patch('chaos_monkey.ChaosMonkey._run_command',
374- autospec=True) as mock:
375- cm.run_chaos('net', 'deny-all', timeout=0)
376- self.assertEqual(0, mock.call_args_list[0][1]['timeout'])
377-
378- def test_run_chaos_raises_for_command_str(self):
379- cm = ChaosMonkey.factory()
380- cm.include_group('all')
381- with patch('utility.check_output', autospec=True):
382- with self.assertRaisesRegexp(
383- NotFound,
384- "Command not found: group: net command_str:foo"):
385- cm.run_chaos('net', 'foo', timeout=0)
386-
387- def test_run_chaos_raises_for_group(self):
388- cm = ChaosMonkey.factory()
389- cm.include_group('all')
390- with patch('utility.check_output', autospec=True):
391- with self.assertRaisesRegexp(
392- NotFound,
393- "Command not found: group: bar command_str:deny-all"):
394- cm.run_chaos('bar', 'deny-all', timeout=0)
395-
396- def test_run_command(self):
397- cm = ChaosMonkey.factory()
398- net = Net()
399- for chaos in net.get_chaos():
400- if chaos.command_str == "deny-state-server":
401- break
402- else:
403- self.fail("'deny-state-server' chaos not found")
404- with patch('utility.check_output', autospec=True) as mock:
405- cm._run_command(chaos, timeout=0)
406- self.assertEqual(mock.mock_calls, [
407- call(['ufw', 'deny', '37017']),
408- call(['ufw', 'allow', 'in', 'to', 'any']),
409- call(['ufw', '--force', 'enable']),
410- call(['ufw', 'disable']),
411- call(['ufw', 'delete', 'allow', 'in', 'to', 'any']),
412- call(['ufw', 'delete', 'deny', '37017']),
413- ])
414-
415- def test_shutdown(self):
416- cm = ChaosMonkey.factory()
417- with patch('utility.check_output', autospec=True) as mock:
418- cm.shutdown()
419- mock.assert_any_call(['ufw', '--force', 'reset'])
420-
421 def test_include_group(self):
422 group = ['net']
423 cm = ChaosMonkey.factory()
424
425=== added file 'tests/test_init.py'
426--- tests/test_init.py 1970-01-01 00:00:00 +0000
427+++ tests/test_init.py 2015-06-05 00:54:01 +0000
428@@ -0,0 +1,87 @@
429+import os
430+from tempfile import NamedTemporaryFile
431+
432+from unittest import TestCase
433+
434+from utils.init import Init
435+
436+__metaclass__ = type
437+
438+
439+class TestInit(TestCase):
440+
441+ def test_upstart(self):
442+ init = Init.upstart()
443+ self.assertIsInstance(init, Init)
444+ self.assertEqual(init.init_path, '/etc/init/chaos-monkey-restart.conf')
445+ cm_dir = get_chaos_monkey_dir()
446+ self.assertEqual(init.init_script_path, os.path.join(
447+ cm_dir, 'scripts', 'chaos-monkey-restart.conf'))
448+ self.assertEqual(init.restart_script_path, os.path.join(
449+ cm_dir, 'scripts', 'restart_chaos_monkey.py'))
450+ self.assertEqual(init.runner_path, os.path.join(
451+ cm_dir, 'runner.py'))
452+
453+ def test_install(self):
454+ init_script_path = os.path.join(
455+ get_chaos_monkey_dir(), 'scripts', 'chaos-monkey-restart.conf')
456+ with open(init_script_path) as f:
457+ conf_content = f.read()
458+ with NamedTemporaryFile() as init_fd:
459+ restart_script_path = '/scripts/restart_chaos_monkey.py'
460+ runner_path = '/path/runner.py'
461+ init = Init(init_path=init_fd.name,
462+ init_script_path=init_script_path,
463+ restart_script_path=restart_script_path,
464+ runner_path=runner_path)
465+ cmd_arg = '--include-command=deny-all workspace'
466+ expire_time = 65537.00
467+ init.install(cmd_arg, expire_time)
468+
469+ conf_content = conf_content.format(
470+ restart_script_path=restart_script_path,
471+ runner_path=runner_path,
472+ expire_time=expire_time,
473+ cmd_arg=cmd_arg)
474+
475+ with open(init_fd.name) as f:
476+ init_content = f.read()
477+ self.assertEqual(init_content, conf_content)
478+
479+ def test_uninstall(self):
480+ with NamedTemporaryFile(delete=False) as init_fd:
481+ init_fd.write('fake')
482+ init_fd.flush()
483+ os.fsync(init_fd)
484+ init = Init(init_path=init_fd.name,
485+ init_script_path='fake',
486+ restart_script_path='fake',
487+ runner_path='fake')
488+ init.uninstall()
489+ self.assertIs(os.path.isfile(init_fd.name), False)
490+
491+ def test_remove_args(self):
492+ cmd_arg = ("--include-command restart-unit --enablement-timeout 1"
493+ " /home/myhome --expire-time 946598401 --restart")
494+ expected = ("--include-command restart-unit --enablement-timeout 1"
495+ " /home/myhome")
496+ result = Init._remove_args(cmd_arg)
497+ self.assertEqual(result, expected)
498+
499+ def test_remove_args_restart_in_middle_of_string(self):
500+ cmd_arg = ("--include-command restart-unit --enablement-timeout 1"
501+ " /home/myhome --restart --expire-time 946598401")
502+ expected = ("--include-command restart-unit --enablement-timeout 1"
503+ " /home/myhome")
504+ result = Init._remove_args(cmd_arg)
505+ self.assertEqual(result, expected)
506+
507+ def test_remove_args_no_expire_time_and_no_restart(self):
508+ cmd_arg = ("--include-command restart-unit --enablement-timeout 1 "
509+ "/home/myhome")
510+ result = Init._remove_args(cmd_arg)
511+ self.assertEqual(result, cmd_arg)
512+
513+
514+def get_chaos_monkey_dir():
515+ return os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
516
517=== modified file 'tests/test_kill.py'
518--- tests/test_kill.py 2015-05-21 23:55:13 +0000
519+++ tests/test_kill.py 2015-06-05 00:54:01 +0000
520@@ -92,4 +92,4 @@
521
522
523 def get_all_kill_commands():
524- return [Kill.jujud_cmd, Kill.mongod_cmd]
525+ return [Kill.jujud_cmd, Kill.mongod_cmd, Kill.restart_cmd]
526
527=== modified file 'tests/test_net.py'
528--- tests/test_net.py 2015-05-28 15:35:16 +0000
529+++ tests/test_net.py 2015-06-05 00:54:01 +0000
530@@ -75,12 +75,6 @@
531 for c in chaos:
532 self.assertEqual('net', c.group)
533
534- def test_shutdown(self):
535- net = Net()
536- with patch('utility.check_output', autospec=True) as mock:
537- net.shutdown()
538- mock.assert_called_once_with(['ufw', '--force', 'reset'])
539-
540 def get_net_chaos(self, cmd):
541 net = Net()
542 for chaos in net.get_chaos():
543
544=== added file 'tests/test_restart_chaos_monkey.py'
545--- tests/test_restart_chaos_monkey.py 1970-01-01 00:00:00 +0000
546+++ tests/test_restart_chaos_monkey.py 2015-06-05 00:54:01 +0000
547@@ -0,0 +1,35 @@
548+from argparse import Namespace
549+
550+from mock import patch
551+from unittest import TestCase
552+
553+from scripts.restart_chaos_monkey import (
554+ parse_args,
555+ restart_chaos_monkey,
556+)
557+
558+__metaclass__ = type
559+
560+
561+class TestRestartChaosMonkey(TestCase):
562+
563+ def test_parse_args(self):
564+ args = parse_args([
565+ '--runner-path', '/path/runner.py',
566+ '--expire-time', '123.00',
567+ '--cmd-arg', '--include-command deny-all workspace'])
568+ self.assertEqual(args, Namespace(
569+ runner_path='/path/runner.py',
570+ expire_time=123.00,
571+ cmd_arg='--include-command deny-all workspace'))
572+
573+ def test_execute_chaos_monkey(self):
574+ args = parse_args([
575+ '--runner-path', '/path/runner.py',
576+ '--expire-time', '123.0',
577+ '--cmd-arg', '--include-command deny-all workspace'])
578+ with patch('subprocess.Popen', autospec=True) as mock:
579+ restart_chaos_monkey(args)
580+ mock.assert_called_once_with(
581+ ['python', '/path/runner.py', '--include-command', 'deny-all',
582+ 'workspace', '--expire-time', '123.0', '--restart'])
583
584=== modified file 'tests/test_runner.py'
585--- tests/test_runner.py 2015-05-22 19:53:05 +0000
586+++ tests/test_runner.py 2015-06-05 00:54:01 +0000
587@@ -6,11 +6,12 @@
588 from StringIO import StringIO
589 from time import time
590
591-from mock import patch
592+from mock import patch, call
593
594 from chaos.kill import Kill
595 from chaos_monkey import ChaosMonkey
596 from chaos_monkey_base import Chaos
597+from chaos.net import Net
598 from runner import (
599 display_all_commands,
600 parse_args,
601@@ -116,14 +117,16 @@
602 with patch('utility.check_output', autospec=True) as mock:
603 with temp_dir() as directory:
604 runner = Runner(directory, ChaosMonkey.factory())
605- runner.random_chaos(run_timeout=1, enablement_timeout=1)
606+ runner.random_chaos(run_timeout=1, enablement_timeout=1,
607+ exclude_command=Kill.restart_cmd)
608 self.assertEqual(mock.called, True)
609
610 def test_random_enablement_zero(self):
611 with patch('utility.check_output', autospec=True) as mock:
612 with temp_dir() as directory:
613 runner = Runner(directory, ChaosMonkey.factory())
614- runner.random_chaos(run_timeout=1, enablement_timeout=0)
615+ runner.random_chaos(run_timeout=1, enablement_timeout=0,
616+ exclude_command=Kill.restart_cmd)
617 self.assertEqual(mock.called, True)
618
619 def test_random_verify_timeout(self):
620@@ -133,44 +136,42 @@
621 with temp_dir() as directory:
622 runner = Runner(directory, ChaosMonkey.factory())
623 runner.random_chaos(run_timeout=run_timeout,
624- enablement_timeout=2)
625+ enablement_timeout=2,
626+ exclude_command=Kill.restart_cmd)
627 end_time = time()
628 self.assertEqual(run_timeout, int(end_time-current_time))
629 self.assertEqual(mock.called, True)
630
631- def test_random_assert_chaos_monkey_methods_called(self):
632- with patch('runner.ChaosMonkey', autospec=True) as cm_mock:
633- with temp_dir() as directory:
634- runner = Runner.factory(directory, ChaosMonkey.factory())
635- runner.random_chaos(run_timeout=1, enablement_timeout=1)
636- cm_mock.factory.return_value.run_random_chaos.assert_called_with(1)
637- cm_mock.factory.return_value.shutdown.assert_called_with()
638+ def test_random_assert_run_command_method_called(self):
639+ with patch('utility.check_output', autospec=True):
640+ with patch('runner.Runner._run_command', autospec=True) as cm_mock:
641+ with temp_dir() as directory:
642+ runner = Runner.factory(directory, ChaosMonkey.factory())
643+ runner.random_chaos(run_timeout=1, enablement_timeout=1)
644+ cm_mock.assert_called_with(runner, 1)
645
646 def test_random_assert_chaos_methods_called(self):
647 net_ctx = patch('chaos.net.Net', autospec=True)
648 kill_ctx = patch('chaos.kill.Kill', autospec=True)
649 with patch('utility.check_output', autospec=True):
650- with patch('chaos_monkey.ChaosMonkey.run_random_chaos',
651- autospec=True):
652+ with patch('runner.Runner._run_command', autospec=True):
653 with net_ctx as net_mock:
654 with kill_ctx as kill_mock:
655 with temp_dir() as directory:
656 runner = Runner(directory, ChaosMonkey.factory())
657- runner.random_chaos(run_timeout=1,
658- enablement_timeout=1)
659- net_mock.factory.return_value.shutdown.assert_called_with()
660- kill_mock.factory.return_value.shutdown.assert_called_with()
661+ runner.random_chaos(
662+ run_timeout=1, enablement_timeout=1)
663 net_mock.factory.return_value.get_chaos.assert_called_with()
664 kill_mock.factory.return_value.get_chaos.assert_called_with()
665
666- def test_random_passes_timeout(self):
667+ def test_random_chaos_passes_timeout(self):
668 with patch('utility.check_output', autospec=True):
669- with patch('chaos_monkey.ChaosMonkey._run_command',
670+ with patch('runner.Runner._run_command',
671 autospec=True) as mock:
672 with temp_dir() as directory:
673 runner = Runner(directory, ChaosMonkey.factory())
674 runner.random_chaos(run_timeout=3, enablement_timeout=2)
675- self.assertEqual(mock.call_args_list[0][1]['timeout'], 2)
676+ self.assertEqual(mock.call_args_list[0][0][1], 2)
677
678 def test_setup_sig_handler_sets_stop_chaos_on_SIGINT(self):
679 with temp_dir() as directory:
680@@ -417,19 +418,18 @@
681 for c in runner.chaos_monkey.chaos))
682
683 def test_filter_commands_gets_options_from_random_chaos(self):
684- with patch('chaos_monkey.ChaosMonkey.run_random_chaos'):
685- with patch('chaos_monkey.ChaosMonkey.shutdown'):
686- with patch('runner.Runner.filter_commands',
687- autospec=True) as f_mock:
688- with temp_dir() as directory:
689- runner = Runner(directory, ChaosMonkey.factory())
690- runner.random_chaos(
691- run_timeout=1,
692- enablement_timeout=1,
693- include_group='net,{}'.format(Kill.group),
694- exclude_group=Kill.group,
695- include_command='deny-all',
696- exclude_command='deny-incoming')
697+ with patch('runner.Runner._run_command', autospec=True):
698+ with patch('runner.Runner.filter_commands',
699+ autospec=True) as f_mock:
700+ with temp_dir() as directory:
701+ runner = Runner(directory, ChaosMonkey.factory())
702+ runner.random_chaos(
703+ run_timeout=1,
704+ enablement_timeout=1,
705+ include_group='net,{}'.format(Kill.group),
706+ exclude_group=Kill.group,
707+ include_command='deny-all',
708+ exclude_command='deny-incoming')
709 expected = {'include_group': 'net,{}'.format(Kill.group),
710 'exclude_group': Kill.group,
711 'include_command': 'deny-all',
712@@ -524,7 +524,7 @@
713 total_timeout=10, log_count=2, include_group=None,
714 exclude_group=None, include_command=None,
715 exclude_command=None, dry_run=False,
716- run_once=False))
717+ run_once=False, restart=False, expire_time=None))
718
719 def test_parse_args_non_default_values(self):
720 args = parse_args(['path',
721@@ -535,14 +535,16 @@
722 '--exclude-group', Kill.group,
723 '--include-command', 'deny-all',
724 '--exclude-command', 'deny-incoming',
725- '--dry-run'])
726+ '--dry-run',
727+ '--restart',
728+ '--expire-time', '111.11'])
729 self.assertEqual(
730 args, Namespace(path='path', enablement_timeout=30,
731 total_timeout=600, log_count=4,
732 include_group='net', exclude_group=Kill.group,
733 include_command='deny-all',
734 exclude_command='deny-incoming', dry_run=True,
735- run_once=False))
736+ run_once=False, restart=True, expire_time=111.11))
737
738 def test_parse_args_non_default_values_set_run_once(self):
739 args = parse_args(['path',
740@@ -560,7 +562,7 @@
741 include_group='net', exclude_group=Kill.group,
742 include_command='deny-all',
743 exclude_command='deny-incoming', dry_run=True,
744- run_once=True))
745+ run_once=True, restart=False, expire_time=None))
746
747 def test_parse_args_error_enablement_greater_than_total_timeout(self):
748 with parse_error(self) as stderr:
749@@ -588,16 +590,13 @@
750
751 def test_random_chaos_run_once(self):
752 cm = ChaosMonkey.factory()
753- with patch('chaos_monkey.ChaosMonkey.run_random_chaos',
754+ with patch('runner.Runner._run_command',
755 autospec=True) as mock:
756- with patch('chaos_monkey.ChaosMonkey.shutdown',
757- autospec=True) as s_mock:
758- with temp_dir() as directory:
759- runner = Runner(directory, cm)
760- runner.random_chaos(
761- run_timeout=2, enablement_timeout=1, run_once=True)
762- mock.assert_called_once_with(cm, 1)
763- s_mock.assert_called_once_with(cm)
764+ with temp_dir() as directory:
765+ runner = Runner(directory, cm)
766+ runner.random_chaos(
767+ run_timeout=2, enablement_timeout=1, run_once=True)
768+ mock.assert_called_once_with(runner, 1)
769
770 def test_list_all_commands(self):
771 cmd = Runner.list_all_commands()
772@@ -623,6 +622,43 @@
773 for item in expected_items:
774 self.assertIn(item, result)
775
776+ def test_run_command(self):
777+ chaos = self._get_chaos_object(Net(), 'deny-state-server')
778+ with patch('utility.check_output', autospec=True) as mock:
779+ with patch(
780+ 'runner.random.choice', autospec=True, return_value=chaos):
781+ with temp_dir() as directory:
782+ runner = Runner(directory, ChaosMonkey.factory())
783+ runner._run_command(enablement_timeout=0)
784+ self.assertEqual(mock.mock_calls, [
785+ call(['ufw', 'deny', '37017']),
786+ call(['ufw', 'allow', 'in', 'to', 'any']),
787+ call(['ufw', '--force', 'enable']),
788+ call(['ufw', 'disable']),
789+ call(['ufw', 'delete', 'allow', 'in', 'to', 'any']),
790+ call(['ufw', 'delete', 'deny', '37017']),
791+ ])
792+
793+ def test_run_command_select_restart_unit(self):
794+ chaos = self._get_chaos_object(Kill(), Kill.restart_cmd)
795+ with patch('utility.check_output', autospec=True) as mock:
796+ with patch(
797+ 'runner.random.choice', autospec=True, return_value=chaos):
798+ with patch('runner.Init', autospec=True) as ri_mock:
799+ with temp_dir() as directory:
800+ runner = Runner(directory, ChaosMonkey.factory())
801+ runner._run_command(enablement_timeout=0)
802+ self.assertEqual(mock.mock_calls, [call(['shutdown', '-r', 'now'])])
803+ ri_mock.upstart.assert_called_once_with()
804+
805+ def _get_chaos_object(self, obj, command_str):
806+ for chaos in obj.get_chaos():
807+ if chaos.command_str == command_str:
808+ break
809+ else:
810+ self.fail("'{}' chaos not found".format(command_str))
811+ return chaos
812+
813
814 def add_fake_group(chaos_monkey):
815 chaos = Chaos(None, None, 'fake_group', 'fake_command_str', 'description')
816
817=== added directory 'utils'
818=== added file 'utils/__init__.py'
819=== added file 'utils/init.py'
820--- utils/init.py 1970-01-01 00:00:00 +0000
821+++ utils/init.py 2015-06-05 00:54:01 +0000
822@@ -0,0 +1,74 @@
823+import errno
824+import logging
825+import os
826+
827+
828+__metaclass__ = type
829+
830+
831+class Init:
832+
833+ def __init__(self, init_path, init_script_path, restart_script_path,
834+ runner_path):
835+ self.init_path = init_path
836+ self.init_script_path = init_script_path
837+ self.restart_script_path = restart_script_path
838+ self.runner_path = runner_path
839+
840+ @classmethod
841+ def upstart(cls):
842+ dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
843+ scripts_dir_path = os.path.join(dir, 'scripts')
844+ init_filename = 'chaos-monkey-restart.conf'
845+ # Full path to upstart conf file.
846+ init_path = os.path.join('/etc/init/', init_filename)
847+ # Full path to upstart script file.
848+ # This file contains upstart conf script.
849+ init_script_path = os.path.join(scripts_dir_path, init_filename)
850+ # Script filename that restarts the Chaos Monkey
851+ restart_script_filename = 'restart_chaos_monkey.py'
852+ # Full path to restart script file.
853+ restart_script_path = os.path.join(scripts_dir_path,
854+ restart_script_filename)
855+ runner_path = os.path.join(dir, 'runner.py')
856+ return cls(
857+ init_path, init_script_path, restart_script_path, runner_path)
858+
859+ def install(self, cmd_arg, expire_time):
860+ cmd_arg = Init._remove_args(cmd_arg=cmd_arg)
861+ with open(self.init_script_path, 'r') as f:
862+ data = f.read().format(
863+ restart_script_path=self.restart_script_path,
864+ runner_path=self.runner_path,
865+ cmd_arg=cmd_arg,
866+ expire_time=expire_time,
867+ )
868+ # Write to /etc/init dir
869+ with open(self.init_path, 'w') as f:
870+ f.write(data)
871+ logging.info("Init script generated:\n cmd: {} \n expire_time: {} \n "
872+ "runner_path: {}".format(cmd_arg, expire_time,
873+ self.runner_path))
874+
875+ @staticmethod
876+ def _remove_args(cmd_arg):
877+ cmd_arg = cmd_arg.replace('--restart ', '').replace('--restart', '')
878+ if '--expire-time' in cmd_arg:
879+ cmd_list = cmd_arg.split()
880+ remove_index = cmd_list.index('--expire-time')
881+ # Delete --expire-time
882+ del cmd_list[remove_index]
883+ # Delete expire-time's argument, which is the next item in the
884+ # list. "remove_index" is going to be the same because an item has
885+ # been removed and the size has been reduced.
886+ del cmd_list[remove_index]
887+ cmd_arg = ' '.join(cmd_list)
888+ return cmd_arg
889+
890+ def uninstall(self):
891+ try:
892+ os.remove(self.init_path)
893+ logging.info("Init script removed from {}".format(self.init_path))
894+ except OSError as e:
895+ if e.errno != errno.ENOENT:
896+ raise

Subscribers

People subscribed via source and target branches