Fix the saltutil.wheel function and add integration tests (#33414)

* Fix the saltutil.wheel function and add integration tests

Fixes #32446

* Provide a more useful message when command is run on non-local minion

salt/modules/saltutil.py

@@ -26,7 +26,7 @@ try:
 except ImportError:
     HAS_ESKY = False
 # pylint: disable=no-name-in-module
-from salt.ext.six import string_types
+import salt.ext.six as six
 from salt.ext.six.moves.urllib.error import URLError
 # pylint: enable=import-error,no-name-in-module
 
@@ -52,7 +52,7 @@ import salt.utils.url
 import salt.transport
 import salt.wheel
 from salt.exceptions import (
-    SaltReqTimeoutError, SaltRenderError, CommandExecutionError
+    SaltReqTimeoutError, SaltRenderError, CommandExecutionError, SaltInvocationError
 )
 
 __proxyenabled__ = ['*']
@@ -88,7 +88,7 @@ def _sync(form, saltenv=None):
     '''
     if saltenv is None:
         saltenv = _get_top_file_envs()
-    if isinstance(saltenv, string_types):
+    if isinstance(saltenv, six.string_types):
         saltenv = saltenv.split(',')
     ret = []
     remote = set()
@@ -989,7 +989,7 @@ def runner(_fun, **kwargs):
     return rclient.cmd(_fun, kwarg=kwargs)
 
 
-def wheel(_fun, **kwargs):
+def wheel(_fun, *args, **kwargs):
     '''
     Execute a wheel module (this function must be run on the master)
 
@@ -997,6 +997,13 @@ def wheel(_fun, **kwargs):
 
     name
         The name of the function to run
+
+    args
+        Any positional arguments to pass to the wheel function. A common example
+        of this would be the ``match`` arg needed for key functions.
+
+        .. versionadded:: v2015.8.11
+
     kwargs
         Any keyword arguments to pass to the wheel function
 
@@ -1004,10 +1011,33 @@ def wheel(_fun, **kwargs):
 
     .. code-block:: bash
 
-        salt '*' saltutil.wheel key.accept match=jerry
+        salt '*' saltutil.wheel key.accept jerry
     '''
-    wclient = salt.wheel.WheelClient(__opts__)
-    return wclient.cmd(_fun, kwarg=kwargs)
+    if __opts__['__role'] == 'minion':
+        master_config = os.path.join(os.path.dirname(__opts__['conf_file']),
+                                     'master')
+        master_opts = salt.config.client_config(master_config)
+        wheel_client = salt.wheel.WheelClient(master_opts)
+    else:
+        wheel_client = salt.wheel.WheelClient(__opts__)
+
+    # The WheelClient cmd needs args, kwargs, and pub_data separated out from
+    # the "normal" kwargs structure, which at this point contains __pub_x keys.
+    pub_data = {}
+    valid_kwargs = {}
+    for key, val in six.iteritems(kwargs):
+        if key.startswith('__'):
+            pub_data[key] = val
+        else:
+            valid_kwargs[key] = val
+
+    try:
+        ret = wheel_client.cmd(_fun, arg=args, pub_data=pub_data, kwarg=valid_kwargs)
+    except SaltInvocationError:
+        raise CommandExecutionError('This command can only be executed on a minion '
+                                    'that is located on the master.')
+
+    return ret
 
 
 # this is the only way I could figure out how to get the REAL file_roots

tests/integration/modules/saltutil.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+'''
+Integration tests for the saltutil module.
+'''
+
+# Import Python libs
+from __future__ import absolute_import
+
+# Import Salt Testing libs
+from salttesting.helpers import ensure_in_syspath
+
+ensure_in_syspath('../../')
+
+# Import Salt libs
+import integration
+
+
+class SaltUtilModuleTest(integration.ModuleCase):
+    '''
+    Testcase for the saltutil execution module
+    '''
+
+    # Tests for the wheel function
+
+    def test_wheel_just_function(self):
+        '''
+        Tests using the saltutil.wheel function when passing only a function.
+        '''
+        ret = self.run_function('saltutil.wheel', ['minions.connected'])
+        self.assertEqual(ret, ['minion', 'sub_minion'])
+
+    def test_wheel_with_arg(self):
+        '''
+        Tests using the saltutil.wheel function when passing a function and an arg.
+        '''
+        ret = self.run_function('saltutil.wheel', ['key.list', 'minion'])
+        self.assertEqual(ret, {})
+
+    def test_wheel_no_arg_raise_error(self):
+        '''
+        Tests using the saltutil.wheel function when passing a function that requires
+        an arg, but one isn't supplied.
+        '''
+        self.assertRaises(TypeError, 'saltutil.wheel', ['key.list'])
+
+    def test_wheel_with_kwarg(self):
+        '''
+        Tests using the saltutil.wheel function when passing a function and a kwarg.
+        This function just generates a key pair, but doesn't do anything with it. We
+        just need this for testing purposes.
+        '''
+        ret = self.run_function('saltutil.wheel', ['key.gen'], keysize=1024)
+        self.assertIn('pub', ret)
+        self.assertIn('priv', ret)
+
+
+if __name__ == '__main__':
+    from integration import run_tests
+    run_tests(SaltUtilModuleTest)