ansible/test/units/executor/test_play_iterator.py

# (c) 2012-2014, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

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 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
from units.mock.path import mock_unfrackpath_noop


class TestPlayIterator(unittest.TestCase):

    def setUp(self):
        pass

    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()


    @patch('ansible.playbook.role.definition.unfrackpath', mock_unfrackpath_noop)
    def test_play_iterator(self):
        #import epdb; epdb.st()
        fake_loader = DictDataLoader({
            "test_play.yml": """
            - hosts: all
              gather_facts: false
              roles:
              - test_role
              pre_tasks:
              - 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"
            """,
            '/etc/ansible/roles/test_role/tasks/main.yml': """
            - debug: msg="this is a role 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)

        mock_var_manager._fact_cache['host00'] = dict()

        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(),
        )

        # lookup up an original task
        target_task = p._entries[0].tasks[0].block[0]
        task_copy = target_task.copy(exclude_parent=True)
        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)
        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')
        # role task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'debug')
        self.assertIsNotNone(task._role)
        # regular play task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        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)
        self.assertEqual(task.action, 'meta')
        # post 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')
        # end of iteration
        (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"
                        - ping:
                      rescue:
                      - block:
                        - block:
                          - block:
                            - block:
                              - debug: msg="this is the rescue task"
                  always:
                  - block:
                    - block:
                      - block:
                        - block:
                          - debug: msg="this is the always 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')
        self.assertEqual(task.args, dict(_raw_params='flush_handlers'))
        # get the first 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 the first task'))
        # 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')
        self.assertEqual(task.args, dict(msg='this is the rescue task'))
        # get the 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 the always task'))
        # implicit meta: flush_handlers
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'meta')
        self.assertEqual(task.args, dict(_raw_params='flush_handlers'))
        # implicit meta: flush_handlers
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'meta')
        self.assertEqual(task.args, dict(_raw_params='flush_handlers'))
        # 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()])