Updating unit tests for PlayIterator

This knowingly introduces a broken test, planning to fix that later.

lib/ansible/executor/play_iterator.py

@@ -58,6 +58,9 @@ class HostState:
         self.always_child_state = None
 
     def __repr__(self):
+        return "HostState(%r)" % self._blocks
+
+    def __str__(self):
         def _run_state_to_string(n):
             states = ["ITERATING_SETUP", "ITERATING_TASKS", "ITERATING_RESCUE", "ITERATING_ALWAYS", "ITERATING_COMPLETE"]
             try:
@@ -90,6 +93,20 @@ class HostState:
             self.always_child_state,
         )
 
+    def __eq__(self, other):
+        if not isinstance(other, HostState):
+            return False
+
+        for attr in (
+            '_blocks', 'cur_block', 'cur_regular_task', 'cur_rescue_task', 'cur_always_task',
+            'cur_role', 'run_state', 'fail_state', 'pending_setup', 'cur_dep_chain',
+            'tasks_child_state', 'rescue_child_state', 'always_child_state'
+            ):
+            if getattr(self, attr) != getattr(other, attr):
+                return False
+
+        return True
+
     def get_current_block(self):
         return self._blocks[self.cur_block]
 
@@ -439,7 +456,7 @@ class PlayIterator:
         the different processes, and not all data structures are preserved. This method
         allows us to find the original task passed into the executor engine.
         '''
-        def _search_block(block, task):
+        def _search_block(block):
             '''
             helper method to check a block's task lists (block/rescue/always)
             for a given task uuid. If a Block is encountered in the place of a
@@ -449,32 +466,32 @@ class PlayIterator:
             for b in (block.block, block.rescue, block.always):
                 for t in b:
                     if isinstance(t, Block):
-                        res = _search_block(t, task)
+                        res = _search_block(t)
                         if res:
                             return res
                     elif t._uuid == task._uuid:
                         return t
             return None
 
-        def _search_state(state, task):
+        def _search_state(state):
             for block in state._blocks:
-                res = _search_block(block, task)
+                res = _search_block(block)
                 if res:
                     return res
             for child_state in (state.tasks_child_state, state.rescue_child_state, state.always_child_state):
                 if child_state is not None:
-                    res = _search_state(child_state, task)
+                    res = _search_state(child_state)
                     if res:
                         return res
             return None
 
         s = self.get_host_state(host)
-        res = _search_state(s, task)
+        res = _search_state(s)
         if res:
             return res
 
         for block in self._play.handlers:
-            res = _search_block(block, task)
+            res = _search_block(block)
             if res:
                 return res
 

lib/ansible/playbook/base.py

@@ -267,6 +267,8 @@ class Base:
         new_me._loader           = self._loader
         new_me._variable_manager = self._variable_manager
 
+        new_me._uuid = self._uuid
+
         # if the ds value was set on the object, copy it to the new copy too
         if hasattr(self, '_ds'):
             new_me._ds = self._ds

test/units/executor/test_play_iterator.py

@@ -23,8 +23,9 @@ from ansible.compat.tests import unittest
 from ansible.compat.tests.mock import patch, MagicMock
 
 from ansible.errors import AnsibleError, AnsibleParserError
-from ansible.executor.play_iterator import PlayIterator
+from ansible.executor.play_iterator import HostState, PlayIterator
 from ansible.playbook import Playbook
+from ansible.playbook.task import Task
 from ansible.playbook.play_context import PlayContext
 
 from units.mock.loader import DictDataLoader
@@ -37,6 +38,23 @@ class TestPlayIterator(unittest.TestCase):
     def tearDown(self):
         pass
 
+    def test_host_state(self):
+        hs = HostState(blocks=[x for x in range(0, 10)])
+        hs.tasks_child_state = HostState(blocks=[0])
+        hs.rescue_child_state = HostState(blocks=[1])
+        hs.always_child_state = HostState(blocks=[2])
+        hs.__repr__()
+        hs.run_state = 100
+        hs.__repr__()
+        hs.fail_state = 15
+        hs.__repr__()
+
+        for i in range(0, 10):
+            hs.cur_block = i
+            self.assertEqual(hs.get_current_block(), i)
+
+        new_hs = hs.copy()
+
     def test_play_iterator(self):
         fake_loader = DictDataLoader({
             "test_play.yml": """
@@ -48,6 +66,18 @@ class TestPlayIterator(unittest.TestCase):
               - debug: msg="this is a pre_task"
               tasks:
               - debug: msg="this is a regular task"
+              - block:
+                - debug: msg="this is a block task"
+                - block:
+                  - debug: msg="this is a sub-block in a block"
+                rescue:
+                - debug: msg="this is a rescue task"
+                - block:
+                  - debug: msg="this is a sub-block in a rescue"
+                always:
+                - debug: msg="this is an always task"
+                - block:
+                  - debug: msg="this is a sub-block in an always"
               post_tasks:
               - debug: msg="this is a post_task"
             """,
@@ -64,10 +94,12 @@ class TestPlayIterator(unittest.TestCase):
 
         hosts = []
         for i in range(0, 10):
-            host  = MagicMock()
+            host = MagicMock()
-            host.get_name.return_value = 'host%02d' % i
+            host.name = host.get_name.return_value = 'host%02d' % i
             hosts.append(host)
 
+        mock_var_manager._fact_cache['host00'] = dict()
+
         inventory = MagicMock()
         inventory.get_hosts.return_value = hosts
         inventory.filter_hosts.return_value = hosts
@@ -82,6 +114,18 @@ class TestPlayIterator(unittest.TestCase):
             all_vars=dict(),
         )
 
+        # lookup up an original task
+        target_task = p._entries[0].tasks[0].block[0]
+        print("the task is: %s (%s)" % (target_task, target_task._uuid))
+        task_copy = target_task.copy(exclude_block=True)
+        print("the copied task is: %s (%s)" % (task_copy, task_copy._uuid))
+        found_task = itr.get_original_task(hosts[0], task_copy)
+        self.assertEqual(target_task, found_task)
+
+        bad_task = Task()
+        found_task = itr.get_original_task(hosts[0], bad_task)
+        self.assertIsNone(found_task)
+
         # pre task
         (host_state, task) = itr.get_next_task_for_host(hosts[0])
         self.assertIsNotNone(task)
@@ -100,6 +144,38 @@ class TestPlayIterator(unittest.TestCase):
         self.assertIsNotNone(task)
         self.assertEqual(task.action, 'debug')
         self.assertIsNone(task._role)
+        # block task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is a block task"))
+        # sub-block task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is a sub-block in a block"))
+        # mark the host failed
+        itr.mark_host_failed(hosts[0])
+        # block rescue task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is a rescue task"))
+        # sub-block rescue task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is a sub-block in a rescue"))
+        # block always task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is an always task"))
+        # sub-block always task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        self.assertEqual(task.args, dict(msg="this is a sub-block in an always"))
         # implicit meta: flush_handlers
         (host_state, task) = itr.get_next_task_for_host(hosts[0])
         self.assertIsNotNone(task)
@@ -116,3 +192,176 @@ class TestPlayIterator(unittest.TestCase):
         (host_state, task) = itr.get_next_task_for_host(hosts[0])
         self.assertIsNone(task)
 
+        # host 0 shouldn't be in the failed hosts, as the error
+        # was handled by a rescue block
+        failed_hosts = itr.get_failed_hosts()
+        self.assertNotIn(hosts[0], failed_hosts)
+
+    def test_play_iterator_nested_blocks(self):
+        fake_loader = DictDataLoader({
+            "test_play.yml": """
+            - hosts: all
+              gather_facts: false
+              tasks:
+              - block:
+                - block:
+                  - block:
+                    - block:
+                      - block:
+                        - debug: msg="this is the first task"
+                    rescue:
+                    - block:
+                      - block:
+                        - block:
+                          - block:
+                            - debug: msg="this is the rescue task"
+                always:
+                - block:
+                  - block:
+                    - block:
+                      - block:
+                        - debug: msg="this is the rescue task"
+            """,
+        })
+
+        mock_var_manager = MagicMock()
+        mock_var_manager._fact_cache = dict()
+        mock_var_manager.get_vars.return_value = dict()
+
+        p = Playbook.load('test_play.yml', loader=fake_loader, variable_manager=mock_var_manager)
+
+        hosts = []
+        for i in range(0, 10):
+            host = MagicMock()
+            host.name = host.get_name.return_value = 'host%02d' % i
+            hosts.append(host)
+
+        inventory = MagicMock()
+        inventory.get_hosts.return_value = hosts
+        inventory.filter_hosts.return_value = hosts
+
+        play_context = PlayContext(play=p._entries[0])
+
+        itr = PlayIterator(
+            inventory=inventory,
+            play=p._entries[0],
+            play_context=play_context,
+            variable_manager=mock_var_manager,
+            all_vars=dict(),
+        )
+
+        # implicit meta: flush_handlers
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'meta')
+        # get the first task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        # fail the host
+        itr.mark_host_failed(hosts[0])
+        # get the resuce task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        # get the always task
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'debug')
+        # implicit meta: flush_handlers
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'meta')
+        # implicit meta: flush_handlers
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNotNone(task)
+        self.assertEqual(task.action, 'meta')
+        # end of iteration
+        (host_state, task) = itr.get_next_task_for_host(hosts[0])
+        self.assertIsNone(task)
+
+    def test_play_iterator_add_tasks(self):
+        fake_loader = DictDataLoader({
+            'test_play.yml': """
+            - hosts: all
+              gather_facts: no
+              tasks:
+              - debug: msg="dummy task"
+            """,
+        })
+
+        mock_var_manager = MagicMock()
+        mock_var_manager._fact_cache = dict()
+        mock_var_manager.get_vars.return_value = dict()
+
+        p = Playbook.load('test_play.yml', loader=fake_loader, variable_manager=mock_var_manager)
+
+        hosts = []
+        for i in range(0, 10):
+            host = MagicMock()
+            host.name = host.get_name.return_value = 'host%02d' % i
+            hosts.append(host)
+
+        inventory = MagicMock()
+        inventory.get_hosts.return_value = hosts
+        inventory.filter_hosts.return_value = hosts
+
+        play_context = PlayContext(play=p._entries[0])
+
+        itr = PlayIterator(
+            inventory=inventory,
+            play=p._entries[0],
+            play_context=play_context,
+            variable_manager=mock_var_manager,
+            all_vars=dict(),
+        )
+
+        # test the high-level add_tasks() method
+        s = HostState(blocks=[0,1,2])
+        itr._insert_tasks_into_state = MagicMock(return_value=s)
+        itr.add_tasks(hosts[0], [3,4,5])
+        self.assertEqual(itr._host_states[hosts[0].name], s)
+
+        # now actually test the lower-level method that does the work
+        itr = PlayIterator(
+            inventory=inventory,
+            play=p._entries[0],
+            play_context=play_context,
+            variable_manager=mock_var_manager,
+            all_vars=dict(),
+        )
+
+        # iterate past first task
+        _, task = itr.get_next_task_for_host(hosts[0])
+        while(task and task.action != 'debug'):
+            _, task = itr.get_next_task_for_host(hosts[0])
+
+        if task is None:
+            raise Exception("iterated past end of play while looking for place to insert tasks")
+
+        # get the current host state and copy it so we can mutate it
+        s = itr.get_host_state(hosts[0])
+        s_copy = s.copy()
+
+        # assert with an empty task list, or if we're in a failed state, we simply return the state as-is
+        res_state = itr._insert_tasks_into_state(s_copy, task_list=[])
+        self.assertEqual(res_state, s_copy)
+
+        s_copy.fail_state = itr.FAILED_TASKS
+        res_state = itr._insert_tasks_into_state(s_copy, task_list=[MagicMock()])
+        self.assertEqual(res_state, s_copy)
+
+        # but if we've failed with a rescue/always block
+        mock_task = MagicMock()
+        s_copy.run_state = itr.ITERATING_RESCUE
+        res_state = itr._insert_tasks_into_state(s_copy, task_list=[mock_task])
+        self.assertEqual(res_state, s_copy)
+        self.assertIn(mock_task, res_state._blocks[res_state.cur_block].rescue)
+        itr._host_states[hosts[0].name] = res_state
+        (next_state, next_task) = itr.get_next_task_for_host(hosts[0], peek=True)
+        self.assertEqual(next_task, mock_task)
+        itr._host_states[hosts[0].name] = s
+
+        # test a regular insertion
+        s_copy = s.copy()
+        res_state = itr._insert_tasks_into_state(s_copy, task_list=[MagicMock()])