diff --git a/salt/modules/win_system.py b/salt/modules/win_system.py index 4952937fc1fd..b866ffeb007f 100644 --- a/salt/modules/win_system.py +++ b/salt/modules/win_system.py @@ -25,6 +25,7 @@ import salt.utils.functools import salt.utils.locales import salt.utils.platform +import salt.utils.winapi from salt.exceptions import CommandExecutionError # Import 3rd-party Libs @@ -527,33 +528,11 @@ def byte_calc(val): else: return '{0:.3f}TB'.format(val / 2**40) - # Connect to WMI - pythoncom.CoInitialize() - conn = wmi.WMI() - # Lookup dicts for Win32_OperatingSystem os_type = {1: 'Work Station', 2: 'Domain Controller', 3: 'Server'} - system = conn.Win32_OperatingSystem()[0] - ret = {'name': get_computer_name(), - 'description': system.Description, - 'install_date': system.InstallDate, - 'last_boot': system.LastBootUpTime, - 'os_manufacturer': system.Manufacturer, - 'os_name': system.Caption, - 'users': system.NumberOfUsers, - 'organization': system.Organization, - 'os_architecture': system.OSArchitecture, - 'primary': system.Primary, - 'os_type': os_type[system.ProductType], - 'registered_user': system.RegisteredUser, - 'system_directory': system.SystemDirectory, - 'system_drive': system.SystemDrive, - 'os_version': system.Version, - 'windows_directory': system.WindowsDirectory} - # lookup dicts for Win32_ComputerSystem domain_role = {0: 'Standalone Workstation', 1: 'Member Workstation', @@ -576,59 +555,82 @@ def byte_calc(val): 6: 'Appliance PC', 7: 'Performance Server', 8: 'Maximum'} - system = conn.Win32_ComputerSystem()[0] - # Get pc_system_type depending on Windows version - if platform.release() in ['Vista', '7', '8']: - # Types for Vista, 7, and 8 - pc_system_type = pc_system_types[system.PCSystemType] - else: - # New types were added with 8.1 and newer - pc_system_types.update({8: 'Slate', 9: 'Maximum'}) - pc_system_type = pc_system_types[system.PCSystemType] - ret.update({ - 'bootup_state': system.BootupState, - 'caption': system.Caption, - 'chassis_bootup_state': warning_states[system.ChassisBootupState], - 'chassis_sku_number': system.ChassisSKUNumber, - 'dns_hostname': system.DNSHostname, - 'domain': system.Domain, - 'domain_role': domain_role[system.DomainRole], - 'hardware_manufacturer': system.Manufacturer, - 'hardware_model': system.Model, - 'network_server_mode_enabled': system.NetworkServerModeEnabled, - 'part_of_domain': system.PartOfDomain, - 'pc_system_type': pc_system_type, - 'power_state': system.PowerState, - 'status': system.Status, - 'system_type': system.SystemType, - 'total_physical_memory': byte_calc(system.TotalPhysicalMemory), - 'total_physical_memory_raw': system.TotalPhysicalMemory, - 'thermal_state': warning_states[system.ThermalState], - 'workgroup': system.Workgroup - }) - # Get processor information - processors = conn.Win32_Processor() - ret['processors'] = 0 - ret['processors_logical'] = 0 - ret['processor_cores'] = 0 - ret['processor_cores_enabled'] = 0 - ret['processor_manufacturer'] = processors[0].Manufacturer - ret['processor_max_clock_speed'] = six.text_type(processors[0].MaxClockSpeed) + 'MHz' - for system in processors: - ret['processors'] += 1 - ret['processors_logical'] += system.NumberOfLogicalProcessors - ret['processor_cores'] += system.NumberOfCores - ret['processor_cores_enabled'] += system.NumberOfEnabledCore - - system = conn.Win32_BIOS()[0] - ret.update({'hardware_serial': system.SerialNumber, - 'bios_manufacturer': system.Manufacturer, - 'bios_version': system.Version, - 'bios_details': system.BIOSVersion, - 'bios_caption': system.Caption, - 'bios_description': system.Description}) - ret['install_date'] = _convert_date_time_string(ret['install_date']) - ret['last_boot'] = _convert_date_time_string(ret['last_boot']) + + # Connect to WMI + with salt.utils.winapi.Com(): + conn = wmi.WMI() + + system = conn.Win32_OperatingSystem()[0] + ret = {'name': get_computer_name(), + 'description': system.Description, + 'install_date': system.InstallDate, + 'last_boot': system.LastBootUpTime, + 'os_manufacturer': system.Manufacturer, + 'os_name': system.Caption, + 'users': system.NumberOfUsers, + 'organization': system.Organization, + 'os_architecture': system.OSArchitecture, + 'primary': system.Primary, + 'os_type': os_type[system.ProductType], + 'registered_user': system.RegisteredUser, + 'system_directory': system.SystemDirectory, + 'system_drive': system.SystemDrive, + 'os_version': system.Version, + 'windows_directory': system.WindowsDirectory} + + system = conn.Win32_ComputerSystem()[0] + # Get pc_system_type depending on Windows version + if platform.release() in ['Vista', '7', '8']: + # Types for Vista, 7, and 8 + pc_system_type = pc_system_types[system.PCSystemType] + else: + # New types were added with 8.1 and newer + pc_system_types.update({8: 'Slate', 9: 'Maximum'}) + pc_system_type = pc_system_types[system.PCSystemType] + ret.update({ + 'bootup_state': system.BootupState, + 'caption': system.Caption, + 'chassis_bootup_state': warning_states[system.ChassisBootupState], + 'chassis_sku_number': system.ChassisSKUNumber, + 'dns_hostname': system.DNSHostname, + 'domain': system.Domain, + 'domain_role': domain_role[system.DomainRole], + 'hardware_manufacturer': system.Manufacturer, + 'hardware_model': system.Model, + 'network_server_mode_enabled': system.NetworkServerModeEnabled, + 'part_of_domain': system.PartOfDomain, + 'pc_system_type': pc_system_type, + 'power_state': system.PowerState, + 'status': system.Status, + 'system_type': system.SystemType, + 'total_physical_memory': byte_calc(system.TotalPhysicalMemory), + 'total_physical_memory_raw': system.TotalPhysicalMemory, + 'thermal_state': warning_states[system.ThermalState], + 'workgroup': system.Workgroup + }) + # Get processor information + processors = conn.Win32_Processor() + ret['processors'] = 0 + ret['processors_logical'] = 0 + ret['processor_cores'] = 0 + ret['processor_cores_enabled'] = 0 + ret['processor_manufacturer'] = processors[0].Manufacturer + ret['processor_max_clock_speed'] = six.text_type(processors[0].MaxClockSpeed) + 'MHz' + for processor in processors: + ret['processors'] += 1 + ret['processors_logical'] += processor.NumberOfLogicalProcessors + ret['processor_cores'] += processor.NumberOfCores + ret['processor_cores_enabled'] += processor.NumberOfEnabledCore + + bios = conn.Win32_BIOS()[0] + ret.update({'hardware_serial': bios.SerialNumber, + 'bios_manufacturer': bios.Manufacturer, + 'bios_version': bios.Version, + 'bios_details': bios.BIOSVersion, + 'bios_caption': bios.Caption, + 'bios_description': bios.Description}) + ret['install_date'] = _convert_date_time_string(ret['install_date']) + ret['last_boot'] = _convert_date_time_string(ret['last_boot']) return ret @@ -692,9 +694,9 @@ def set_hostname(hostname): salt 'minion-id' system.set_hostname newhostname ''' - pythoncom.CoInitialize() - conn = wmi.WMI() - comp = conn.Win32_ComputerSystem()[0] + with salt.utils.winapi.Com(): + conn = wmi.WMI() + comp = conn.Win32_ComputerSystem()[0] return comp.Rename(Name=hostname) @@ -970,13 +972,13 @@ def get_domain_workgroup(): salt 'minion-id' system.get_domain_workgroup ''' - pythoncom.CoInitialize() - conn = wmi.WMI() - for computer in conn.Win32_ComputerSystem(): - if computer.PartOfDomain: - return {'Domain': computer.Domain} - else: - return {'Workgroup': computer.Workgroup} + with salt.utils.winapi.Com(): + conn = wmi.WMI() + for computer in conn.Win32_ComputerSystem(): + if computer.PartOfDomain: + return {'Domain': computer.Domain} + else: + return {'Workgroup': computer.Workgroup} def set_domain_workgroup(workgroup): @@ -998,14 +1000,14 @@ def set_domain_workgroup(workgroup): workgroup = _to_unicode(workgroup) # Initialize COM - pythoncom.CoInitialize() + with salt.utils.winapi.Com(): + # Grab the first Win32_ComputerSystem object from wmi + conn = wmi.WMI() + comp = conn.Win32_ComputerSystem()[0] - # Grab the first Win32_ComputerSystem object from wmi - conn = wmi.WMI() - comp = conn.Win32_ComputerSystem()[0] + # Now we can join the new workgroup + res = comp.JoinDomainOrWorkgroup(Name=workgroup.upper()) - # Now we can join the new workgroup - res = comp.JoinDomainOrWorkgroup(Name=workgroup.upper()) return True if not res[0] else False diff --git a/tests/unit/modules/test_win_system.py b/tests/unit/modules/test_win_system.py index 198f2cbae30b..894706493e67 100644 --- a/tests/unit/modules/test_win_system.py +++ b/tests/unit/modules/test_win_system.py @@ -14,21 +14,152 @@ from tests.support.mock import ( MagicMock, patch, + Mock, NO_MOCK, NO_MOCK_REASON ) # Import Salt Libs import salt.modules.win_system as win_system +import salt.utils.stringutils + +try: + import wmi + HAS_WMI = True +except ImportError: + HAS_WMI = False + + +class MockWMI_ComputerSystem(object): + ''' + Mock WMI Win32_ComputerSystem Class + ''' + BootupState = 'Normal boot' + Caption = 'SALT SERVER' + ChassisBootupState = 3 + ChassisSKUNumber = '3.14159' + DNSHostname = 'SALT SERVER' + Domain = 'WORKGROUP' + DomainRole = 2 + Manufacturer = 'Dell Inc.' + Model = 'Dell 2980' + NetworkServerModeEnabled = True + PartOfDomain = False + PCSystemType = 4 + PowerState = 0 + Status = 'OK' + SystemType = 'x64-based PC' + TotalPhysicalMemory = 17078214656 + ThermalState = 3 + Workgroup = 'WORKGROUP' + + def __init__(self): + pass + + @staticmethod + def Rename(Name): + return Name == Name + + @staticmethod + def JoinDomainOrWorkgroup(Name): + return [0] + + +class MockWMI_OperatingSystem(object): + ''' + Mock WMI Win32_OperatingSystem Class + ''' + Description = 'Because salt goes EVERYWHERE' + InstallDate = '20110211131800' + LastBootUpTime = '19620612120000' + Manufacturer = 'Python' + Caption = 'Salty' + NumberOfUsers = 7530000000 + Organization = 'SaltStack' + OSArchitecture = 'Windows' + Primary = True + ProductType = 3 + RegisteredUser = 'thatch@saltstack.com' + SystemDirectory = 'C:\\Windows\\System32' + SystemDrive = 'C:\\' + Version = '10.0.17763' + WindowsDirectory = 'C:\\Windows' + + def __init__(self): + pass + + +class MockWMI_Processor(object): + ''' + Mock WMI Win32_Processor Class + ''' + Manufacturer = 'Intel' + MaxClockSpeed = 2301 + NumberOfLogicalProcessors = 8 + NumberOfCores = 4 + NumberOfEnabledCore = 4 + + def __init__(self): + pass + + +class MockWMI_BIOS(object): + ''' + Mock WMI Win32_BIOS Class + ''' + SerialNumber = 'SALTY2011' + Manufacturer = 'Dell Inc.' + Version = 'DELL - 10283849' + Caption = 'A12' + BIOSVersion = [Version, Caption, 'ASUS - 3948D'] + Description = Caption + + def __init__(self): + pass + + +class Mockwinapi(object): + ''' + Mock winapi class + ''' + def __init__(self): + pass + + class winapi(object): + ''' + Mock winapi class + ''' + def __init__(self): + pass + + @staticmethod + def Com(): + ''' + Mock Com method + ''' + return True @skipIf(NO_MOCK, NO_MOCK_REASON) +@skipIf(not HAS_WMI, 'WMI only available on Windows') class WinSystemTestCase(TestCase, LoaderModuleMockMixin): ''' Test cases for salt.modules.win_system ''' def setup_loader_modules(self): modules_globals = {} + # wmi and pythoncom modules are platform specific... + mock_pythoncom = types.ModuleType( + salt.utils.stringutils.to_str('pythoncom') + ) + sys_modules_patcher = patch.dict('sys.modules', + {'pythoncom': mock_pythoncom}) + sys_modules_patcher.start() + self.addCleanup(sys_modules_patcher.stop) + self.WMI = Mock() + self.addCleanup(delattr, self, 'WMI') + modules_globals['wmi'] = wmi + if win_system.HAS_WIN32NET_MODS is False: win32api = types.ModuleType( str('win32api') # future lint: disable=blacklisted-function @@ -38,6 +169,7 @@ def setup_loader_modules(self): now.day, now.hour, now.minute, now.second, now.microsecond]) modules_globals['win32api'] = win32api + return {win_system: modules_globals} def test_halt(self): @@ -181,10 +313,9 @@ def test_set_computer_desc(self): with patch.dict(win_system.__salt__, {'cmd.run': mock}): mock = MagicMock(return_value="Salt's comp") with patch.object(win_system, 'get_computer_desc', mock): - self.assertDictEqual(win_system.set_computer_desc( - "Salt's comp" - ), - {'Computer Description': "Salt's comp"}) + self.assertDictEqual( + win_system.set_computer_desc("Salt's comp"), + {'Computer Description': "Salt's comp"}) @skipIf(not win_system.HAS_WIN32NET_MODS, 'Missing win32 libraries') def test_get_computer_desc(self): @@ -277,17 +408,39 @@ def test_set_hostname(self): ''' Test setting a new hostname ''' - cmd_run_mock = MagicMock(return_value="Method execution successful.") - get_hostname = MagicMock(return_value="MINION") - with patch.dict(win_system.__salt__, {'cmd.run': cmd_run_mock}): - with patch.object(win_system, 'get_hostname', get_hostname): - win_system.set_hostname("NEW") + with patch('salt.utils', Mockwinapi), \ + patch('salt.utils.winapi.Com', MagicMock()), \ + patch.object(self.WMI, 'Win32_ComputerSystem', + return_value=[MockWMI_ComputerSystem()]), \ + patch.object(wmi, 'WMI', Mock(return_value=self.WMI)): + self.assertTrue(win_system.set_hostname("NEW")) + + def test_get_domain_workgroup(self): + ''' + Test get_domain_workgroup + ''' + with patch('salt.utils', Mockwinapi), \ + patch.object(wmi, 'WMI', Mock(return_value=self.WMI)), \ + patch('salt.utils.winapi.Com', MagicMock()), \ + patch.object(self.WMI, 'Win32_ComputerSystem', + return_value=[MockWMI_ComputerSystem()]): + self.assertDictEqual(win_system.get_domain_workgroup(), + {'Workgroup': 'WORKGROUP'}) - cmd_run_mock.assert_called_once_with(cmd="wmic computersystem where name='MINION' call rename name='NEW'") + def test_set_domain_workgroup(self): + ''' + Test set_domain_workgroup + ''' + with patch('salt.utils', Mockwinapi), \ + patch.object(wmi, 'WMI', Mock(return_value=self.WMI)), \ + patch('salt.utils.winapi.Com', MagicMock()), \ + patch.object(self.WMI, 'Win32_ComputerSystem', + return_value=[MockWMI_ComputerSystem()]): + self.assertTrue(win_system.set_domain_workgroup('test')) def test_get_hostname(self): ''' - Test setting a new hostname + Test getting a new hostname ''' cmd_run_mock = MagicMock(return_value="MINION") with patch.dict(win_system.__salt__, {'cmd.run': cmd_run_mock}): @@ -314,7 +467,19 @@ def test_get_system_info(self): 'thermal_state', 'total_physical_memory', 'total_physical_memory_raw', 'users', 'windows_directory', 'workgroup'] - ret = win_system.get_system_info() + with patch('salt.utils', Mockwinapi), \ + patch('salt.utils.winapi.Com', MagicMock()), \ + patch.object(self.WMI, 'Win32_OperatingSystem', + return_value=[MockWMI_OperatingSystem()]), \ + patch.object(self.WMI, 'Win32_ComputerSystem', + return_value=[MockWMI_ComputerSystem()]), \ + patch.object(self.WMI, 'Win32_Processor', + return_value=[MockWMI_Processor(), + MockWMI_Processor()]), \ + patch.object(self.WMI, 'Win32_BIOS', + return_value=[MockWMI_BIOS()]), \ + patch.object(wmi, 'WMI', Mock(return_value=self.WMI)): + ret = win_system.get_system_info() # Make sure all the fields are in the return for field in fields: self.assertIn(field, ret)