saltstack/salt
1
0
Fork 0
mirror of https://github.com/saltstack/salt.git synced 2025-04-17 10:10:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Merge pull request #46655 from dwoz/pyobjects-46350

Browse source 
Fixing cleanUp method to restore environment
This commit is contained in:
Nicole Thomas 2018-03-29 14:31:47 -04:00 • committed by GitHub
commit 395b7f8fdc
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 45 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
salt/states/environ.py
View file

@ -20,6 +20,15 @@ def __virtual__():
return True
def _norm_key(key):
'''
Normalize windows environment keys
'''
if utils.is_windows():
return key.upper()
return key
def setenv(name,
value,
false_unsets=False,
@ -124,12 +133,11 @@ def setenv(name,
permanent_hive = 'HKLM'
permanent_key = r'SYSTEM\CurrentControlSet\Control\Session Manager\Environment'
out = __salt__['reg.read_value'](permanent_hive, permanent_key, key)
out = __salt__['reg.read_value'](permanent_hive, permanent_key, _norm_key(key))
return out['success'] is True
else:
return False
if current_environ.get(key, None) is None and not key_exists():
if current_environ.get(_norm_key(key), None) is None and not key_exists():
# The key does not exist in environment
if false_unsets is not True:
# This key will be added with value ''
@ -138,13 +146,13 @@ def setenv(name,
# The key exists.
if false_unsets is not True:
# Check to see if the value will change
if current_environ.get(key, None) != '':
if current_environ.get(_norm_key(key), None) != '':
# This key value will change to ''
ret['changes'].update({key: ''})
else:
# We're going to delete the key
ret['changes'].update({key: None})
elif current_environ.get(key, '') == val:
elif current_environ.get(_norm_key(key), '') == val:
already_set.append(key)
else:
ret['changes'].update({key: val})

45
tests/unit/states/test_environ.py
View file

@ -6,7 +6,7 @@ import os
# Import Salt Testing libs
from tests.support.mixins import LoaderModuleMockMixin
from tests.support.unit import TestCase
from tests.support.unit import TestCase, skipIf
from tests.support.mock import (
MagicMock,
patch
@ -14,6 +14,8 @@ from tests.support.mock import (
# Import salt libs
import salt.states.environ as envstate
import salt.modules.environ as envmodule
import salt.modules.reg
import salt.utils
class TestEnvironState(TestCase, LoaderModuleMockMixin):
@ -22,17 +24,21 @@ class TestEnvironState(TestCase, LoaderModuleMockMixin):
loader_globals = {
'__env__': 'base',
'__opts__': {'test': False},
'__salt__': {'environ.setenv': envmodule.setenv}
'__salt__': {
'environ.setenv': envmodule.setenv,
'reg.read_value': salt.modules.reg.read_value,
}
}
return {envstate: loader_globals, envmodule: loader_globals}
def setUp(self):
original_environ = os.environ.copy()
os.environ = {'INITIAL': 'initial'}
patcher = patch.dict(os.environ, {'INITIAL': 'initial'}, clear=True)
patcher.start()
def reset_environ(original_environ):
os.environ = original_environ
self.addCleanup(reset_environ, original_environ)
def reset_environ(patcher):
patcher.stop()
self.addCleanup(reset_environ, patcher)
def test_setenv(self):
'''test that a subsequent calls of setenv changes nothing'''
@ -47,9 +53,10 @@ class TestEnvironState(TestCase, LoaderModuleMockMixin):
ret = envstate.setenv('test', 'other')
self.assertEqual(ret['changes'], {})
@skipIf(not salt.utils.is_windows(), 'Windows only')
def test_setenv_permanent(self):
with patch.dict(envmodule.__salt__, {'reg.set_value': MagicMock(), 'reg.delete_value': MagicMock()}), \
patch('salt.utils.is_windows', MagicMock(return_value=True)):
'''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')
@ -82,14 +89,20 @@ class TestEnvironState(TestCase, LoaderModuleMockMixin):
'''test that ``clear_all`` option sets other values to '' '''
ret = envstate.setenv('test', 'value', clear_all=True)
self.assertEqual(ret['changes'], {'test': 'value', 'INITIAL': ''})
self.assertEqual(envstate.os.environ, {'test': 'value', 'INITIAL': ''})
if salt.utils.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})
self.assertEqual(envstate.os.environ, {'test': 'value'})
if salt.utils.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'''
@ -100,12 +113,18 @@ class TestEnvironState(TestCase, LoaderModuleMockMixin):
ret = envstate.setenv(
'notimportant', {'test': False, 'foo': 'baz'}, false_unsets=True)
self.assertEqual(ret['changes'], {'test': None, 'foo': 'baz'})
self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'foo': 'baz'})
if salt.utils.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'})
self.assertEqual(envstate.os.environ, {'INITIAL': 'initial', 'foo': 'bax', 'test': ''})
if salt.utils.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'''