path: root/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/>.

from __future__ import annotations

import unittest
from unittest.mock import patch, MagicMock

from ansible.executor.play_iterator import HostState, PlayIterator, IteratingStates, FailedStates
from ansible.playbook import Playbook
from ansible.playbook.play_context import PlayContext
from ansible.plugins.loader import init_plugin_loader

from units.mock.loader import DictDataLoader
from units.mock.path import mock_unfrackpath_noop


class TestPlayIterator(unittest.TestCase):

    def test_host_state(self):
        hs = HostState(blocks=list(range(0, 10)))
        hs.tasks_child_state = HostState(blocks=[0])
        hs.rescue_child_state = HostState(blocks=[1])
        hs.always_child_state = HostState(blocks=[2])
        repr(hs)
        hs.run_state = 100
        repr(hs)
        hs.fail_state = 15
        repr(hs)

        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):
        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': """
            - name: role task
              debug: msg="this is a role task"
            - block:
              - name: role block task
                debug: msg="inside block in role"
              always:
              - name: role always task
                debug: msg="always task in block in role"
            - name: role include_tasks
              include_tasks: foo.yml
            - name: role task after include
              debug: msg="after include in role"
            - block:
              - name: starting role nested block 1
                debug:
              - block:
                - name: role nested block 1 task 1
                  debug:
                - name: role nested block 1 task 2
                  debug:
                - name: role nested block 1 task 3
                  debug:
              - name: end of role nested block 1
                debug:
              - name: starting role nested block 2
                debug:
              - block:
                - name: role nested block 2 task 1
                  debug:
                - name: role nested block 2 task 2
                  debug:
                - name: role nested block 2 task 3
                  debug:
              - name: end of role nested block 2
                debug:
            """,
            '/etc/ansible/roles/test_role/tasks/foo.yml': """
            - name: role included task
              debug: msg="this is task in an include from a role"
            """
        })

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

        # pre task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'debug')
        # role task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'debug')
        self.assertEqual(task.name, "role task")
        self.assertIsNotNone(task._role)
        # role block task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role block task")
        self.assertIsNotNone(task._role)
        # role block always task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role always task")
        self.assertIsNotNone(task._role)
        # role include_tasks
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'include_tasks')
        self.assertEqual(task.name, "role include_tasks")
        self.assertIsNotNone(task._role)
        # role task after include
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role task after include")
        self.assertIsNotNone(task._role)
        # role nested block tasks
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "starting role nested block 1")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 1 task 1")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 1 task 2")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 1 task 3")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "end of role nested block 1")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "starting role nested block 2")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 2 task 1")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 2 task 2")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "role nested block 2 task 3")
        self.assertIsNotNone(task._role)
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.name, "end of role nested block 2")
        self.assertIsNotNone(task._role)
        # implicit meta: role_complete
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'meta')
        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"))
        # post task
        (host_state, task) = itr.get_next_task_for_host(hosts[0])
        self.assertIsNotNone(task)
        self.assertEqual(task.action, 'debug')
        # 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):
        init_plugin_loader()
        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(),
        )

        # 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 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 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'))
        # 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], [MagicMock(), MagicMock(), MagicMock()])
        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
        dummy, task = itr.get_next_task_for_host(hosts[0])
        while (task and task.action != 'debug'):
            dummy, task = itr.get_next_task_for_host(hosts[0])

        self.assertIsNotNone(task, '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 = FailedStates.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 = IteratingStates.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.set_state_for_host(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.set_state_for_host(hosts[0].name, s)

        # test a regular insertion
        s_copy = s.copy()
        res_state = itr._insert_tasks_into_state(s_copy, task_list=[MagicMock()])