Merge lp:~sseman/juju-chaos-monkey/replay-commands into lp:juju-chaos-monkey

[Seman]

This branch adds support to replay Chaos Monkey commands from a structured (YAML) data file. It handles a reboot requests as follows:
  - Write the remaining Chaos Monkey commands to a temporary file.
  - After a reboot, read the temporary file and continue the run.
  - Delete the temporary file.

Here is output from using the --replay option:

$ cat /home/ubuntu/cmd_list.log
- [deny-incoming, 2]
- [restart-unit, 2]
- [deny-all, 2]

$runner.py --replay /home/ubuntu/cmd_list.log /home/ubuntu/cm

$cat results.log
2015-06-10 01:02:41 INFO Chaos Monkey started in /home/ubuntu/cm
2015-06-10 01:02:41 INFO Replaying commands from /home/ubuntu/cm/cmd_list.log
2015-06-10 01:02:41 INFO Deny all incoming network traffic except ssh.
2015-06-10 01:02:45 INFO Restart the unit.
2015-06-10 01:02:45 INFO Init script generated:
 cmd: --replay /home/ubuntu/cm/cmd_list.log /home/ubuntu/cm
 expire_time: 1433898167.64
 runner_path: /home/ubuntu/cm/runner.py
2015-06-10 01:02:45 INFO Chaos Monkey stopped.

2015-06-10 01:03:08 INFO Chaos Monkey restarted after a reboot in /home/ubuntu/cm
2015-06-10 01:03:08 INFO Init script removed from /etc/init/chaos-monkey-restart.conf
2015-06-10 01:03:08 INFO Replaying commands from /home/ubuntu/cm/cmd_list.log
2015-06-10 01:03:08 INFO Deny all incoming and outgoing network traffic except ssh.
2015-06-10 01:03:15 INFO Chaos Monkey stopped.

[John George (jog)]

This looks good, thanks for the example output in the description, it helps a lot while reviewing.

If you passed args returned from parse_args() to Runner.factory and setup the constructor to accept each value, with some minimal default, then all the instance data can be stored in one place. The function definitions can then be smaller. Rather than passing in args to each function, self can be referenced. Changing the data storage location is not required to land this branch but we've talked about a later branch that will change argument passing for the random_chaos() function. Please strongly consider this proposed change for that branch.

review: Approve

=== modified file 'chaos_monkey.py'
--- chaos_monkey.py	2015-06-04 17:49:00 +0000
+++ chaos_monkey.py	2015-06-10 01:20:20 +0000
@@ -18,14 +18,6 @@
         all_chaos, factory_obj = ChaosMonkey.get_all_chaos()
         return cls([], factory_obj)
 
-    @property
-    def command_tag(self):
-        return ":CHAOS_CMD:"
-
-    @property
-    def description_tag(self):
-        return ":CHAOS_DSCR:"
-
     @staticmethod
     def get_all_chaos():
         all_chaos = []
@@ -85,3 +77,6 @@
             if item.command_str == command_str:
                 return item
         return None
+
+    def reset_command_selection(self):
+        self.chaos = []

=== modified file 'runner.py'
--- runner.py	2015-06-08 23:51:53 +0000
+++ runner.py	2015-06-10 01:20:20 +0000
@@ -10,6 +10,8 @@
     sleep
 )
 
+import yaml
+
 from chaos.kill import Kill
 from chaos_monkey import ChaosMonkey
 from utility import (
@@ -35,6 +37,7 @@
         self.chaos_monkey = chaos_monkey
         self.expire_time = None
         self.cmd_log_name = cmd_log_name
+        self.replay_filename_ext = '.part'
 
     @classmethod
     def factory(cls, workspace, log_count=1, dry_run=False):
@@ -131,6 +134,7 @@
                         include_command=None, exclude_command=None):
         all_groups = ChaosMonkey.get_all_groups()
         all_commands = ChaosMonkey.get_all_commands()
+        self.chaos_monkey.reset_command_selection()
 
         # If any groups and any commands are not included, assume the intent
         #  is to include all groups and all commands.
@@ -151,6 +155,42 @@
                 exclude_command, all_commands)
             self.chaos_monkey.exclude_command(exclude_command)
 
+    def replay_commands(self, args):
+        """Replay Chaos Monkey commands from a file."""
+        commands = self.get_command_list(args)
+        while commands:
+            command = commands.pop()
+            command_str = command[0]
+            enablement_timeout = command[1]
+            if command_str == Kill.restart_cmd and commands:
+                # Save the commands to a temporary file before a reboot.
+                self.save_command_list(commands, args)
+            self.random_chaos(
+                run_timeout=enablement_timeout,
+                enablement_timeout=enablement_timeout,
+                include_command=command_str)
+            if command_str == Kill.restart_cmd:
+                break
+
+    def get_command_list(self, args):
+        """Get the command list from a file."""
+        file_path = (args.replay + self.replay_filename_ext
+                     if args.restart else args.replay)
+        with open(file_path) as f:
+            commands = yaml.load(f.read())
+        commands.reverse()
+        # If it is a restart, remove the temporary file.
+        if args.restart:
+            os.remove(file_path)
+        return commands
+
+    def save_command_list(self, commands, args):
+        """Before a shutdown and restart request, this method is called
+        to save the command list to a temporary file."""
+        file_path = args.replay + self.replay_filename_ext
+        with open(file_path, 'w') as f:
+            f.write(yaml.dump(commands))
+
     @staticmethod
     def _validate(sub_string, all_list):
         sub_list = split_arg_string(sub_string)
@@ -245,6 +285,9 @@
     parser.add_argument(
         '-ep', '--expire-time', type=float,
         help='Chaos Monkey expire time (UNIX timestamp).', default=None)
+    parser.add_argument(
+        '-rp', '--replay', metavar='FULL-FILE-PATH',
+        help='Replay Chaos Monkey commands from a file.', default=None)
     args = parser.parse_args(argv)
 
     if args.run_once and args.total_timeout:
@@ -262,8 +305,12 @@
     if args.enablement_timeout < 0:
         parser.error("Invalid enablement-timeout value: timeout must be "
                      "zero or greater.")
+    if args.replay and not os.path.isabs(args.replay):
+            parser.error("Please provide an absolute file path to the replay "
+                         "argument: {}".format(args.replay))
     return args
 
+
 if __name__ == '__main__':
     args = parse_args()
     runner = Runner.factory(workspace=args.path, log_count=args.log_count,
@@ -278,15 +325,19 @@
 
     runner.acquire_lock(restart=args.restart)
     try:
-        runner.random_chaos(
-            run_timeout=args.total_timeout,
-            enablement_timeout=args.enablement_timeout,
-            include_group=args.include_group,
-            exclude_group=args.exclude_group,
-            include_command=args.include_command,
-            exclude_command=args.exclude_command,
-            run_once=args.run_once,
-            expire_time=args.expire_time)
+        if args.replay:
+            logging.info('Replaying commands from {}'.format(args.replay))
+            runner.replay_commands(args=args)
+        else:
+            runner.random_chaos(
+                run_timeout=args.total_timeout,
+                enablement_timeout=args.enablement_timeout,
+                include_group=args.include_group,
+                exclude_group=args.exclude_group,
+                include_command=args.include_command,
+                exclude_command=args.exclude_command,
+                run_once=args.run_once,
+                expire_time=args.expire_time)
     except Exception as e:
         logging.error('{} ({})'.format(e, type(e).__name__))
         sys.exit(1)

=== modified file 'tests/test_runner.py'
--- tests/test_runner.py	2015-06-08 23:51:53 +0000
+++ tests/test_runner.py	2015-06-10 01:20:20 +0000
@@ -4,9 +4,11 @@
 import signal
 import subprocess
 from StringIO import StringIO
+from tempfile import NamedTemporaryFile
 from time import time
 
 from mock import patch, call
+import yaml
 
 from chaos.kill import Kill
 from chaos_monkey import ChaosMonkey
@@ -530,7 +532,8 @@
                             total_timeout=10, log_count=2, include_group=None,
                             exclude_group=None, include_command=None,
                             exclude_command=None, dry_run=False,
-                            run_once=False, restart=False, expire_time=None))
+                            run_once=False, restart=False, expire_time=None,
+                            replay=None))
 
     def test_parse_args_non_default_values(self):
         args = parse_args(['path',
@@ -543,14 +546,16 @@
                            '--exclude-command', 'deny-incoming',
                            '--dry-run',
                            '--restart',
-                           '--expire-time', '111.11'])
+                           '--expire-time', '111.11',
+                           '--replay', '/path/to/foo'])
         self.assertEqual(
             args, Namespace(path='path', enablement_timeout=30,
                             total_timeout=600, log_count=4,
                             include_group='net', exclude_group=Kill.group,
                             include_command='deny-all',
                             exclude_command='deny-incoming', dry_run=True,
-                            run_once=False, restart=True, expire_time=111.11))
+                            run_once=False, restart=True, expire_time=111.11,
+                            replay='/path/to/foo'))
 
     def test_parse_args_non_default_values_set_run_once(self):
         args = parse_args(['path',
@@ -568,7 +573,8 @@
                             include_group='net', exclude_group=Kill.group,
                             include_command='deny-all',
                             exclude_command='deny-incoming', dry_run=True,
-                            run_once=True, restart=False, expire_time=None))
+                            run_once=True, restart=False, expire_time=None,
+                            replay=None))
 
     def test_parse_args_error_enablement_greater_than_total_timeout(self):
         with parse_error(self) as stderr:
@@ -636,14 +642,7 @@
                 with temp_dir() as directory:
                     runner = Runner(directory, ChaosMonkey.factory())
                     runner._run_command(enablement_timeout=0)
-        self.assertEqual(mock.mock_calls, [
-            call(['ufw', 'deny', '37017']),
-            call(['ufw', 'allow', 'in', 'to', 'any']),
-            call(['ufw', '--force', 'enable']),
-            call(['ufw', 'disable']),
-            call(['ufw', 'delete', 'allow', 'in', 'to', 'any']),
-            call(['ufw', 'delete', 'deny', '37017']),
-        ])
+        self.assertEqual(mock.mock_calls, self._deny_port_call_list())
 
     def test_run_command_select_restart_unit(self):
         chaos = self._get_chaos_object(Kill(), Kill.restart_cmd)
@@ -657,6 +656,80 @@
         self.assertEqual(mock.mock_calls, [call(['shutdown', '-r', 'now'])])
         ri_mock.upstart.assert_called_once_with()
 
+    def test_replay_commands(self):
+        with patch('utility.check_output', autospec=True) as mock:
+            with temp_dir() as directory:
+                runner = Runner(directory, ChaosMonkey.factory())
+                with NamedTemporaryFile() as temp_file:
+                    self._write_command_list_to_file(temp_file)
+                    args = Namespace(replay=temp_file.name, restart=False)
+                    runner.replay_commands(args)
+        expected = self._deny_port_call_list()
+        expected.extend(self._deny_port_call_list('17017'))
+        self.assertEqual(mock.mock_calls, expected)
+
+    def test_replay_commands_with_restart_command(self):
+        commands = "- [restart-unit, 1]\n- [deny-api-server, 1]\n"
+        with patch('utility.check_output', autospec=True) as mock:
+            with patch('runner.Init.install'):
+                with temp_dir() as directory:
+                    runner = Runner(directory, ChaosMonkey.factory())
+                    with NamedTemporaryFile() as temp_file:
+                        self._write_command_list_to_file(
+                            temp_file, data=commands)
+                        args = Namespace(replay=temp_file.name, restart=False)
+                        runner.replay_commands(args)
+                        # Verify a temporary file is created because there
+                        # is a restart-unit command in the list.
+                        self.assertIs(os.path.isfile(
+                            temp_file.name + runner.replay_filename_ext), True)
+                        args = Namespace(replay=temp_file.name, restart=True)
+                        file_content = runner.get_command_list(args)
+                        # Verify the temporary files is deleted.
+                        self.assertIsNot(os.path.isfile(
+                            temp_file.name + runner.replay_filename_ext), True)
+        self.assertEqual(mock.mock_calls, [call(['shutdown', '-r', 'now'])])
+        self.assertEqual(file_content, [yaml.load(commands)[1]])
+
+    def test_replay_commands_after_reboot(self):
+        with patch('utility.check_output', autospec=True) as mock:
+            with temp_dir() as directory:
+                runner = Runner(directory, ChaosMonkey.factory())
+                temp_file = NamedTemporaryFile(
+                    suffix=runner.replay_filename_ext, delete=False)
+                self._write_command_list_to_file(temp_file)
+                args = Namespace(
+                    replay=temp_file.name.split('.')[0], restart=True)
+                runner.replay_commands(args)
+                # Verify replay_commands() has deleted the temp file.
+                self.assertIsNot(os.path.isfile(temp_file.name), True)
+        expected = self._deny_port_call_list()
+        expected.extend(self._deny_port_call_list('17017'))
+        self.assertEqual(mock.mock_calls, expected)
+
+    def test_get_command_list(self):
+        with temp_dir() as directory:
+            runner = Runner(directory, ChaosMonkey.factory())
+            with NamedTemporaryFile() as temp_file:
+                self._write_command_list_to_file(temp_file)
+                args = Namespace(replay=temp_file.name, restart=False)
+                commands = runner.get_command_list(args)
+        expected = [['deny-state-server', 1], ['deny-api-server', 1]]
+        self.assertItemsEqual(commands, expected)
+
+    def test_save_replay_command_list(self):
+        commands = [['deny-state-server', 1], ['deny-api-server', 1]]
+        with temp_dir() as directory:
+            runner = Runner(directory, ChaosMonkey.factory())
+            with NamedTemporaryFile(
+                    suffix=runner.replay_filename_ext) as temp_file:
+                args = Namespace(replay=temp_file.name.split('.')[0],
+                                 restart=False)
+                runner.save_command_list(commands, args)
+                file_content = temp_file.read()
+        expected = yaml.dump(commands)
+        self.assertItemsEqual(file_content, expected)
+
     def _get_chaos_object(self, obj, command_str):
         for chaos in obj.get_chaos():
             if chaos.command_str == command_str:
@@ -665,6 +738,22 @@
             self.fail("'{}' chaos not found".format(command_str))
         return chaos
 
+    def _write_command_list_to_file(self, fd, data=None):
+        data = ("- [deny-state-server, 1]\n- [deny-api-server, 1]\n"
+                if not data else data)
+        fd.write(data)
+        fd.flush()
+        return data
+
+    def _deny_port_call_list(self, port='37017'):
+        return [
+            call(['ufw', 'deny', port]),
+            call(['ufw', 'allow', 'in', 'to', 'any']),
+            call(['ufw', '--force', 'enable']),
+            call(['ufw', 'disable']),
+            call(['ufw', 'delete', 'allow', 'in', 'to', 'any']),
+            call(['ufw', 'delete', 'deny', port])]
+
 
 def add_fake_group(chaos_monkey):
     chaos = Chaos(None, None, 'fake_group', 'fake_command_str', 'description')