Add all attached drives to the disks grain

salt/grains/disks.py

@@ -157,14 +157,13 @@ def _windows_disks():
 
     namespace = r'\\root\microsoft\windows\storage'
     path = 'MSFT_PhysicalDisk'
-    where = '(MediaType=3 or MediaType=4)'
     get = 'DeviceID,MediaType'
 
     ret = {'disks': [], 'SSDs': []}
 
     cmdret = __salt__['cmd.run_all'](
-        '{0} /namespace:{1} path {2} where {3} get {4} /format:table'.format(
-            wmic, namespace, path, where, get))
+        '{0} /namespace:{1} path {2} get {3} /format:table'.format(
+            wmic, namespace, path, get))
 
     if cmdret['retcode'] != 0:
         log.trace('Disk grain does not support this version of Windows')
@@ -181,10 +180,12 @@ def _windows_disks():
             elif mediatype == '4':
                 log.trace('Device %s reports itself as an SSD', device)
                 ret['SSDs'].append(device)
+                ret['disks'].append(device)
+            elif mediatype == '5':
+                log.trace('Device %s reports itself as an SCM', device)
+                ret['disks'].append(device)
             else:
-                log.trace(
-                    'Unable to identify device %s as an SSD or HDD. It does '
-                    'not report 3 or 4', device
-                )
+                log.trace('Device %s reports itself as Unspecified', device)
+                ret['disks'].append(device)
 
     return ret

tests/unit/grains/test_disks.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+'''
+    :codeauthor: :email:`Shane Lee <slee@saltstack.com>`
+'''
+# Import Python libs
+from __future__ import absolute_import, print_function, unicode_literals
+import textwrap
+
+# Import Salt Testing Libs
+from tests.support.mixins import LoaderModuleMockMixin
+from tests.support.unit import TestCase, skipIf
+from tests.support.mock import (
+    patch,
+    MagicMock,
+    NO_MOCK,
+    NO_MOCK_REASON
+)
+
+# Import Salt Libs
+import salt.grains.disks as disks
+
+
+@skipIf(NO_MOCK, NO_MOCK_REASON)
+class IscsiGrainsTestCase(TestCase, LoaderModuleMockMixin):
+    '''
+    Test cases for _windows_disks grains
+    '''
+    def setup_loader_modules(self):
+        return {
+            disks: {
+                '__salt__': {},
+            },
+        }
+
+    def test__windows_disks(self):
+        '''
+        Test grains._windows_disks, normal return
+        Should return a populated dictionary
+        '''
+        mock_which = MagicMock(return_value='C:\\Windows\\System32\\wbem\\WMIC.exe')
+        wmic_result = textwrap.dedent('''
+            DeviceId  MediaType
+            0         4
+            1         0
+            2         3
+            3         5
+        ''')
+        mock_run_all = MagicMock(return_value={'stdout': wmic_result,
+                                               'retcode': 0})
+
+        with patch('salt.utils.path.which', mock_which), \
+                patch.dict(disks.__salt__, {'cmd.run_all': mock_run_all}):
+            result = disks._windows_disks()
+            expected = {
+                'SSDs': ['\\\\.\\PhysicalDrive0'],
+                'disks': [
+                    '\\\\.\\PhysicalDrive0',
+                    '\\\\.\\PhysicalDrive1',
+                    '\\\\.\\PhysicalDrive2',
+                    '\\\\.\\PhysicalDrive3']}
+            self.assertDictEqual(result, expected)
+            cmd = ' '.join([
+                'C:\\Windows\\System32\\wbem\\WMIC.exe',
+                '/namespace:\\\\root\\microsoft\\windows\\storage',
+                'path',
+                'MSFT_PhysicalDisk',
+                'get',
+                'DeviceID,MediaType',
+                '/format:table'
+            ])
+            mock_run_all.assert_called_once_with(cmd)
+
+    def test__windows_disks_retcode(self):
+        '''
+        Test grains._windows_disks, retcode 1
+        Should return empty lists
+        '''
+        mock_which = MagicMock(return_value='C:\\Windows\\System32\\wbem\\WMIC.exe')
+        mock_run_all = MagicMock(return_value={'stdout': '',
+                                               'retcode': 1})
+        with patch('salt.utils.path.which', mock_which), \
+                patch.dict(disks.__salt__, {'cmd.run_all': mock_run_all}):
+            result = disks._windows_disks()
+            expected = {
+                'SSDs': [],
+                'disks': []}
+            self.assertDictEqual(result, expected)