1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6import sys
7import unittest
8
9from knack.commands import CLICommandsLoader, CommandGroup
10from knack.arguments import CLIArgumentType, CLICommandArgument, ArgumentsContext
11from tests.util import MockContext
12
13
14def _dictContainsSubset(expected, actual):
15    """Checks whether actual is a superset of expected.
16       Helper for deprecated assertDictContainsSubset"""
17    missing = False
18    mismatched = False
19    for key, value in expected.items():
20        if key not in actual:
21            missing = True
22        elif value != actual[key]:
23            mismatched = True
24    return False if missing or mismatched else True
25
26
27class TestCommandRegistration(unittest.TestCase):
28
29    def setUp(self):
30        self.mock_ctx = MockContext()
31
32    @staticmethod
33    def sample_command_handler(group_name, resource_name, opt_param=None, expand=None):
34        """
35        The operation to get a virtual machine.
36
37        :param group_name: The name of the group.
38        :type group_name: str
39        :param resource_name: The name of the resource.
40        :type resource_name: str
41        :param opt_param: Used to verify reflection correctly
42        identifies optional params.
43        :type opt_param: object
44        :param expand: The expand expression to apply on the operation.
45        :type expand: str
46        :param dict custom_headers: headers that will be added to the request
47        :param boolean raw: returns the direct response alongside the
48            deserialized response
49        """
50        pass
51
52    @staticmethod
53    def sample_command_handler2(group_name, resource_name, opt_param=None, expand=None, custom_headers=None,
54                                raw=False, **operation_config):
55        pass
56
57    def _set_command_name(self, command):
58        self.mock_ctx.invocation.data['command_string'] = command
59        return command
60
61    def test_register_cli_argument(self):
62        cl = CLICommandsLoader(self.mock_ctx)
63        command_name = self._set_command_name('test register sample-command')
64        with CommandGroup(cl, 'test register', '{}#{{}}'.format(__name__)) as g:
65            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
66                                                       TestCommandRegistration.sample_command_handler.__name__))
67        with ArgumentsContext(cl, command_name) as ac:
68            ac.argument('resource_name', CLIArgumentType(
69                options_list=('--wonky-name', '-n'), metavar='RNAME', help='Completely WONKY name...',
70                required=False
71            ))
72        cl.load_arguments(command_name)
73        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
74        command_metadata = cl.command_table[command_name]
75        self.assertEqual(len(command_metadata.arguments), 4, 'We expected exactly 4 arguments')
76        some_expected_arguments = {
77            'group_name': CLIArgumentType(dest='group_name', required=True),
78            'resource_name': CLIArgumentType(dest='resource_name', required=False),
79        }
80        for probe in some_expected_arguments:
81            existing = next(arg for arg in command_metadata.arguments if arg == probe)
82            contains_subset = _dictContainsSubset(some_expected_arguments[existing].settings,
83                                                  command_metadata.arguments[existing].options)
84            self.assertTrue(contains_subset)
85        self.assertEqual(command_metadata.arguments['resource_name'].options_list, ('--wonky-name', '-n'))
86
87    def test_register_command_custom_excluded_params(self):
88        command_name = self._set_command_name('test sample-command')
89        ep = ['self', 'raw', 'custom_headers', 'operation_config', 'content_version', 'kwargs', 'client']
90        cl = CLICommandsLoader(self.mock_ctx, excluded_command_handler_args=ep)
91        with CommandGroup(cl, 'test', '{}#{{}}'.format(__name__)) as g:
92            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
93                                                       TestCommandRegistration.sample_command_handler2.__name__))
94        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
95        cl.load_arguments(command_name)
96        command_metadata = cl.command_table[command_name]
97        self.assertEqual(len(command_metadata.arguments), 4, 'We expected exactly 4 arguments')
98        self.assertIn(command_name, cl.command_table)
99
100    def test_register_command(self):
101        cl = CLICommandsLoader(self.mock_ctx)
102        command_name = self._set_command_name('test register sample-command')
103        with CommandGroup(cl, 'test register', '{}#{{}}'.format(__name__)) as g:
104            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
105                                                       TestCommandRegistration.sample_command_handler.__name__))
106
107        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
108        cl.load_arguments(command_name)
109        command_metadata = cl.command_table[command_name]
110        self.assertEqual(len(command_metadata.arguments), 4, 'We expected exactly 4 arguments')
111        some_expected_arguments = {
112            'group_name': CLIArgumentType(dest='group_name',
113                                          required=True,
114                                          help='The name of the group.'),
115            'resource_name': CLIArgumentType(dest='resource_name',
116                                             required=True,
117                                             help='The name of the resource.'),
118            'opt_param': CLIArgumentType(required=False,
119                                         help='Used to verify reflection correctly identifies optional params.'),
120            'expand': CLIArgumentType(required=False,
121                                      help='The expand expression to apply on the operation.')
122        }
123
124        for probe in some_expected_arguments:
125            existing = next(arg for arg in command_metadata.arguments if arg == probe)
126            contains_subset = _dictContainsSubset(some_expected_arguments[existing].settings,
127                                                  command_metadata.arguments[existing].options)
128            self.assertTrue(contains_subset)
129        self.assertEqual(command_metadata.arguments['resource_name'].options_list, ['--resource-name'])
130
131    def test_register_command_group_with_no_group_name(self):
132        cl = CLICommandsLoader(self.mock_ctx)
133        command_name = self._set_command_name('sample-command')
134        with CommandGroup(cl, None, '{}#{{}}'.format(__name__)) as g:
135            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
136                                                       TestCommandRegistration.sample_command_handler.__name__))
137
138        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
139        self.assertIn(command_name, cl.command_table)
140
141    def test_register_command_confirmation_bool(self):
142        cl = CLICommandsLoader(self.mock_ctx)
143        command_name = self._set_command_name('test sample-command')
144        with CommandGroup(cl, 'test', '{}#{{}}'.format(__name__)) as g:
145            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
146                                                       TestCommandRegistration.sample_command_handler.__name__),
147                      confirmation=True)
148        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
149        cl.load_arguments(command_name)
150        command_metadata = cl.command_table[command_name]
151        self.assertIn('yes', command_metadata.arguments)
152        self.assertEqual(command_metadata.arguments['yes'].type.settings['action'], 'store_true')
153        self.assertIs(command_metadata.confirmation, True)
154
155    def test_register_command_confirmation_callable(self):
156        cl = CLICommandsLoader(self.mock_ctx)
157
158        def confirm_callable(_):
159            pass
160        command_name = self._set_command_name('test sample-command')
161        with CommandGroup(cl, 'test', '{}#{{}}'.format(__name__)) as g:
162            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
163                                                       TestCommandRegistration.sample_command_handler.__name__),
164                      confirmation=confirm_callable)
165        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
166        cl.load_arguments(command_name)
167        command_metadata = cl.command_table[command_name]
168        self.assertIn('yes', command_metadata.arguments)
169        self.assertEqual(command_metadata.arguments['yes'].type.settings['action'], 'store_true')
170        self.assertIs(command_metadata.confirmation, confirm_callable)
171
172    def test_register_cli_argument_with_overrides(self):
173        cl = CLICommandsLoader(self.mock_ctx)
174        base_type = CLIArgumentType(options_list=['--foo', '-f'], metavar='FOO', help='help1')
175        derived_type = CLIArgumentType(base_type=base_type, help='help2')
176        with CommandGroup(cl, 'test', '{}#{{}}'.format(__name__)) as g:
177            g.command('sample-get', '{}.{}'.format(TestCommandRegistration.__name__,
178                                                   TestCommandRegistration.sample_command_handler.__name__))
179            g.command('command sample-get-1', '{}.{}'.format(TestCommandRegistration.__name__,
180                                                             TestCommandRegistration.sample_command_handler.__name__))
181            g.command('command sample-get-2', '{}.{}'.format(TestCommandRegistration.__name__,
182                                                             TestCommandRegistration.sample_command_handler.__name__))
183        self.assertEqual(len(cl.command_table), 3, 'We expect exactly three commands in the command table')
184
185        def test_with_command(command, target_value):
186            self._set_command_name(command)
187            with ArgumentsContext(cl, 'test') as c:
188                c.argument('resource_name', base_type)
189            with ArgumentsContext(cl, 'test command') as c:
190                c.argument('resource_name', derived_type)
191            with ArgumentsContext(cl, 'test command sample-get-2') as c:
192                c.argument('resource_name', derived_type, help='help3')
193            cl.load_arguments(command)
194            command1 = cl.command_table[command].arguments['resource_name']
195            self.assertEqual(command1.options['help'], target_value)
196
197        test_with_command('test sample-get', 'help1')
198        test_with_command('test command sample-get-1', 'help2')
199        test_with_command('test command sample-get-2', 'help3')
200
201    def test_register_extra_cli_argument(self):
202        cl = CLICommandsLoader(self.mock_ctx)
203        command_name = self._set_command_name('test register sample-command')
204        with CommandGroup(cl, 'test register', '{}#{{}}'.format(__name__)) as g:
205            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
206                                                       TestCommandRegistration.sample_command_handler.__name__))
207        with ArgumentsContext(cl, command_name) as ac:
208            ac.extra('added_param', options_list=('--added-param',),
209                     metavar='ADDED', help='Just added this right now!', required=True)
210        cl.load_arguments(command_name)
211        self.assertEqual(len(cl.command_table), 1, 'We expect exactly one command in the command table')
212        command_metadata = cl.command_table[command_name]
213        self.assertEqual(len(command_metadata.arguments), 5, 'We expected exactly 5 arguments')
214
215        some_expected_arguments = {
216            'added_param': CLIArgumentType(dest='added_param', required=True)
217        }
218
219        for probe in some_expected_arguments:
220            existing = next(arg for arg in command_metadata.arguments if arg == probe)
221            contains_subset = _dictContainsSubset(some_expected_arguments[existing].settings,
222                                                  command_metadata.arguments[existing].options)
223            self.assertTrue(contains_subset)
224
225    def test_register_ignore_cli_argument(self):
226        cl = CLICommandsLoader(self.mock_ctx)
227        command_name = self._set_command_name('test register sample-command')
228        self.mock_ctx.invocation.data['command_string'] = command_name
229        with CommandGroup(cl, 'test register', '{}#{{}}'.format(__name__)) as g:
230            g.command('sample-command', '{}.{}'.format(TestCommandRegistration.__name__,
231                                                       TestCommandRegistration.sample_command_handler.__name__))
232        with ArgumentsContext(cl, 'test register') as ac:
233            ac.argument('resource_name', options_list=['--this'])
234        with ArgumentsContext(cl, 'test register sample-command') as ac:
235            ac.ignore('resource_name')
236            ac.argument('opt_param', options_list=['--this'])
237        cl.load_arguments(command_name)
238        self.assertNotEqual(cl.command_table[command_name].arguments['resource_name'].options_list,
239                            cl.command_table[command_name].arguments['opt_param'].options_list,
240                            "Name conflict in options list")
241
242    def test_command_build_argument_help_text(self):
243        def sample_sdk_method_with_weird_docstring(param_a, param_b, param_c):  # pylint: disable=unused-argument
244            """
245            An operation with nothing good.
246
247            :param dict param_a:
248            :param param_b: The name
249            of
250            nothing.
251            :param param_c: The name
252            of
253
254            nothing2.
255            """
256            pass
257
258        cl = CLICommandsLoader(self.mock_ctx)
259        command_name = self._set_command_name('test command foo')
260        setattr(sys.modules[__name__], sample_sdk_method_with_weird_docstring.__name__,
261                sample_sdk_method_with_weird_docstring)
262        with CommandGroup(cl, 'test command', '{}#{{}}'.format(__name__)) as g:
263            g.command('foo', sample_sdk_method_with_weird_docstring.__name__)
264        cl.load_arguments(command_name)
265        command_metadata = cl.command_table[command_name]
266        self.assertEqual(len(command_metadata.arguments), 3, 'We expected exactly 3 arguments')
267        some_expected_arguments = {
268            'param_a': CLIArgumentType(dest='param_a', required=True, help=''),
269            'param_b': CLIArgumentType(dest='param_b', required=True, help='The name of nothing.'),
270            'param_c': CLIArgumentType(dest='param_c', required=True, help='The name of nothing2.')
271        }
272
273        for probe in some_expected_arguments:
274            existing = next(arg for arg in command_metadata.arguments if arg == probe)
275            contains_subset = _dictContainsSubset(some_expected_arguments[existing].settings,
276                                                  command_metadata.arguments[existing].options)
277            self.assertTrue(contains_subset)
278
279    def test_override_existing_option_string(self):
280        arg = CLIArgumentType(options_list=('--funky', '-f'))
281        updated_options_list = ('--something-else', '-s')
282        arg.update(options_list=updated_options_list, validator=lambda: (), completer=lambda: ())
283        self.assertEqual(arg.settings['options_list'], updated_options_list)
284        self.assertIsNotNone(arg.settings['validator'])
285        self.assertIsNotNone(arg.settings['completer'])
286
287    def test_dont_override_existing_option_string(self):
288        existing_options_list = ('--something-else', '-s')
289        arg = CLIArgumentType(options_list=existing_options_list)
290        arg.update()
291        self.assertEqual(arg.settings['options_list'], existing_options_list)
292
293    def test_override_remove_validator(self):
294        existing_options_list = ('--something-else', '-s')
295        arg = CLIArgumentType(options_list=existing_options_list,
296                              validator=lambda *args, **kwargs: ())
297        arg.update(validator=None)
298        self.assertIsNone(arg.settings['validator'])
299
300    def test_override_using_register_cli_argument(self):
301        def sample_sdk_method(param_a):  # pylint: disable=unused-argument
302            pass
303
304        def test_validator_completer():
305            pass
306
307        cl = CLICommandsLoader(self.mock_ctx)
308        command_name = self._set_command_name('override_using_register_cli_argument foo')
309        setattr(sys.modules[__name__], sample_sdk_method.__name__, sample_sdk_method)
310        with CommandGroup(cl, 'override_using_register_cli_argument', '{}#{{}}'.format(__name__)) as g:
311            g.command('foo', sample_sdk_method.__name__)
312        with ArgumentsContext(cl, 'override_using_register_cli_argument') as ac:
313            ac.argument('param_a',
314                        options_list=('--overridden', '-r'),
315                        validator=test_validator_completer,
316                        completer=test_validator_completer,
317                        required=False)
318        cl.load_arguments(command_name)
319
320        command_metadata = cl.command_table[command_name]
321        self.assertEqual(len(command_metadata.arguments), 1, 'We expected exactly 1 arguments')
322
323        actual_arg = command_metadata.arguments['param_a']
324        self.assertEqual(actual_arg.options_list, ('--overridden', '-r'))
325        self.assertEqual(actual_arg.validator, test_validator_completer)
326        self.assertEqual(actual_arg.completer, test_validator_completer)
327        self.assertFalse(actual_arg.options['required'])
328
329    def test_override_argtype_with_argtype(self):
330        existing_options_list = ('--default', '-d')
331        arg = CLIArgumentType(options_list=existing_options_list, validator=None, completer='base',
332                              help='base', required=True)
333        overriding_argtype = CLIArgumentType(options_list=('--overridden',), validator='overridden',
334                                             completer=None, overrides=arg, help='overridden',
335                                             required=CLIArgumentType.REMOVE)
336        self.assertEqual(overriding_argtype.settings['validator'], 'overridden')
337        self.assertIsNone(overriding_argtype.settings['completer'])
338        self.assertEqual(overriding_argtype.settings['options_list'], ('--overridden',))
339        self.assertEqual(overriding_argtype.settings['help'], 'overridden')
340        self.assertEqual(overriding_argtype.settings['required'], CLIArgumentType.REMOVE)
341
342        cmd_arg = CLICommandArgument(dest='whatever', argtype=overriding_argtype,
343                                     help=CLIArgumentType.REMOVE)
344        self.assertNotIn('required', cmd_arg.options)
345        self.assertNotIn('help', cmd_arg.options)
346
347    def test_cli_ctx_type_error(self):
348        with self.assertRaises(TypeError):
349            CLICommandsLoader(cli_ctx=object())
350
351
352if __name__ == '__main__':
353    unittest.main()
354