salt/tests/unit/states/test_environ.py

# -*- coding: utf-8 -*-

# Import python libs
from __future__ import absolute_import, print_function, unicode_literals
import os

# Import Salt Testing libs
from tests.support.mixins import LoaderModuleMockMixin
from tests.support.unit import TestCase, skipIf
from tests.support.mock import (
    MagicMock,
    patch
)
# Import salt libs
import salt.states.environ as envstate
import salt.modules.environ as envmodule
import salt.modules.reg
import salt.utils.platform


class TestEnvironState(TestCase, LoaderModuleMockMixin):

    def setup_loader_modules(self):
        loader_globals = {
            '__env__': 'base',
            '__opts__': {'test': False},
            '__salt__': {
                'environ.setenv': envmodule.setenv,
                'reg.read_value': salt.modules.reg.read_value,
            }
        }
        return {envstate: loader_globals, envmodule: loader_globals}

    def setUp(self):
        patcher = patch.dict(os.environ, {'INITIAL': 'initial'}, clear=True)
        patcher.start()

        def reset_environ(patcher):
            patcher.stop()

        self.addCleanup(reset_environ, patcher)

    def test_setenv(self):
        '''
        test that a subsequent calls of setenv changes nothing
        '''
        ret = envstate.setenv('test', 'value')
        self.assertEqual(ret['changes'], {'test': 'value'})

        ret = envstate.setenv('test', 'other')
        self.assertEqual(ret['changes'], {'test': 'other'})

        # once again with the same value
        ret = envstate.setenv('test', 'other')
        self.assertEqual(ret['changes'], {})

    @skipIf(not salt.utils.platform.is_windows(), 'Windows only')
    def test_setenv_permanent(self):
        '''
        test that we can set perminent environment variables (requires pywin32)
        '''
        with patch.dict(envmodule.__salt__, {'reg.set_value': MagicMock(), 'reg.delete_value': MagicMock()}):
            ret = envstate.setenv('test', 'value', permanent=True)
            self.assertEqual(ret['changes'], {'test': 'value'})
            envmodule.__salt__['reg.set_value'].assert_called_with("HKCU", "Environment", 'test', 'value')

            ret = envstate.setenv('test', False, false_unsets=True, permanent=True)
            self.assertEqual(ret['changes'], {'test': None})
            envmodule.__salt__['reg.delete_value'].assert_called_with("HKCU", "Environment", 'test')

    def test_setenv_dict(self):
        '''
        test that setenv can be invoked with dict
        '''
        ret = envstate.setenv('notimportant', {'test': 'value'})
        self.assertEqual(ret['changes'], {'test': 'value'})

    def test_setenv_int(self):
        '''
        test that setenv can not be invoked with int
        (actually it's anything other than strings and dict)
        '''
        ret = envstate.setenv('test', 1)
        self.assertEqual(ret['result'], False)

    def test_setenv_unset(self):
        '''
        test that ``false_unsets`` option removes variable from environment
        '''
        ret = envstate.setenv('test', 'value')
        self.assertEqual(ret['changes'], {'test': 'value'})

        ret = envstate.setenv('notimportant', {'test': False}, false_unsets=True)
        self.assertEqual(ret['changes'], {'test': None})
        self.assertEqual(envstate.os.environ, {'INITIAL': 'initial'})

    def test_setenv_clearall(self):
        '''
        test that ``clear_all`` option sets other values to ''
        '''
        ret = envstate.setenv('test', 'value', clear_all=True)
        self.assertEqual(ret['changes'], {'test': 'value', 'INITIAL': ''})
        if salt.utils.platform.is_windows():
            self.assertEqual(envstate.os.environ, {'TEST': 'value', 'INITIAL': ''})
        else:
            self.assertEqual(envstate.os.environ, {'test': 'value', 'INITIAL': ''})

    def test_setenv_clearall_with_unset(self):
        '''
        test that ``clear_all`` option combined with ``false_unsets``
        unsets other values from environment
        '''
        ret = envstate.setenv('test', 'value', false_unsets=True, clear_all=True)
        self.assertEqual(ret['changes'], {'test': 'value', 'INITIAL': None})
        if salt.utils.platform.is_windows():
            self.assertEqual(envstate.os.environ, {'TEST': 'value'})
        else:
            self.assertEqual(envstate.os.environ, {'test': 'value'})

    def test_setenv_unset_multi(self):
        '''
        test basically same things that above tests but with multiple values passed
        '''
        ret = envstate.setenv('notimportant', {'foo': 'bar'})
        self.assertEqual(ret['changes'], {'foo': 'bar'})

        with patch.dict(envstate.__salt__, {'reg.read_value': MagicMock()}):
            ret = envstate.setenv(
                'notimportant', {'test': False, 'foo': 'baz'}, false_unsets=True)
        self.assertEqual(ret['changes'], {'test': None, 'foo': 'baz'})
        if salt.utils.platform.is_windows():
            self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'FOO': 'baz'})
        else:
            self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'foo': 'baz'})

        with patch.dict(envstate.__salt__, {'reg.read_value': MagicMock()}):
            ret = envstate.setenv('notimportant', {'test': False, 'foo': 'bax'})
        self.assertEqual(ret['changes'], {'test': '', 'foo': 'bax'})
        if salt.utils.platform.is_windows():
            self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'FOO': 'bax', 'TEST': ''})
        else:
            self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'foo': 'bax', 'test': ''})

    def test_setenv_test_mode(self):
        '''
        test that imitating action returns good values
        '''
        with patch.dict(envstate.__opts__, {'test': True}):
            ret = envstate.setenv('test', 'value')
            self.assertEqual(ret['changes'], {'test': 'value'})
            ret = envstate.setenv('INITIAL', 'initial')
            self.assertEqual(ret['changes'], {})