Add --key-name and --key-unset option for server rebuild API.

Change-Id: I6d4793a8e961080ea1d6d414cef8d6bbed0c53e7
Story: 2002609
Task: 22228
Signed-off-by: Fan Zhang <zh.f@outlook.com>

2 files changed, 116 insertions, 0 deletions

diff --git a/openstackclient/compute/v2/server.py b/openstackclient/compute/v2/server.py
index 67b4140a..5723dae3 100644
--- a/openstackclient/compute/v2/server.py
+++ b/openstackclient/compute/v2/server.py
@@ -1529,6 +1529,22 @@ class RebuildServer(command.ShowOne):
             action='store_true',
             help=_('Wait for rebuild to complete'),
         )
+        key_group = parser.add_mutually_exclusive_group()
+        key_group.add_argument(
+            '--key-name',
+            metavar='<key-name>',
+            help=_("Set the key name of key pair on the rebuilt instance."
+                   " Cannot be specified with the '--key-unset' option."
+                   " (Supported by API versions '2.54' - '2.latest')"),
+        )
+        key_group.add_argument(
+            '--key-unset',
+            action='store_true',
+            default=False,
+            help=_("Unset the key name of key pair on the rebuilt instance."
+                   " Cannot be specified with the '--key-name' option."
+                   " (Supported by API versions '2.54' - '2.latest')"),
+        )
         return parser
 
     def take_action(self, parsed_args):
@@ -1553,6 +1569,16 @@ class RebuildServer(command.ShowOne):
         if parsed_args.property:
             kwargs['meta'] = parsed_args.property
 
+        if parsed_args.key_name or parsed_args.key_unset:
+            if compute_client.api_version < api_versions.APIVersion('2.54'):
+                msg = _('--os-compute-api-version 2.54 or later is required')
+                raise exceptions.CommandError(msg)
+
+        if parsed_args.key_unset:
+            kwargs['key_name'] = None
+        if parsed_args.key_name:
+            kwargs['key_name'] = parsed_args.key_name
+
         server = server.rebuild(image, parsed_args.password, **kwargs)
         if parsed_args.wait:
             if utils.wait_for_status(
diff --git a/openstackclient/tests/unit/compute/v2/test_server.py b/openstackclient/tests/unit/compute/v2/test_server.py
index e938564b..9adc3cc6 100644
--- a/openstackclient/tests/unit/compute/v2/test_server.py
+++ b/openstackclient/tests/unit/compute/v2/test_server.py
@@ -2517,6 +2517,96 @@ class TestServerRebuild(TestServer):
         self.server.rebuild.assert_called_with(
             self.image, None, meta=expected_property)
 
+    def test_rebuild_with_keypair_name(self):
+        self.server.key_name = 'mykey'
+        arglist = [
+            self.server.id,
+            '--key-name', self.server.key_name,
+        ]
+        verifylist = [
+            ('server', self.server.id),
+            ('key_name', self.server.key_name)
+        ]
+        parsed_args = self.check_parser(self.cmd, arglist, verifylist)
+
+        self.app.client_manager.compute.api_version = 2.54
+        with mock.patch.object(api_versions,
+                               'APIVersion',
+                               return_value=2.54):
+            self.cmd.take_action(parsed_args)
+            args = (
+                self.image,
+                None,
+            )
+            kwargs = dict(
+                key_name=self.server.key_name,
+            )
+            self.servers_mock.get.assert_called_with(self.server.id)
+            self.images_mock.get.assert_called_with(self.image.id)
+            self.server.rebuild.assert_called_with(*args, **kwargs)
+
+    def test_rebuild_with_keypair_name_older_version(self):
+        self.server.key_name = 'mykey'
+        arglist = [
+            self.server.id,
+            '--key-name', self.server.key_name,
+        ]
+        verifylist = [
+            ('server', self.server.id),
+            ('key_name', self.server.key_name)
+        ]
+        parsed_args = self.check_parser(self.cmd, arglist, verifylist)
+
+        self.app.client_manager.compute.api_version = 2.53
+        with mock.patch.object(api_versions,
+                               'APIVersion',
+                               return_value=2.54):
+            self.assertRaises(exceptions.CommandError,
+                              self.cmd.take_action,
+                              parsed_args)
+
+    def test_rebuild_with_keypair_unset(self):
+        self.server.key_name = 'mykey'
+        arglist = [
+            self.server.id,
+            '--key-unset',
+        ]
+        verifylist = [
+            ('server', self.server.id),
+        ]
+        parsed_args = self.check_parser(self.cmd, arglist, verifylist)
+
+        self.app.client_manager.compute.api_version = 2.54
+        with mock.patch.object(api_versions,
+                               'APIVersion',
+                               return_value=2.54):
+            self.cmd.take_action(parsed_args)
+            args = (
+                self.image,
+                None,
+            )
+            kwargs = dict(
+                key_name=None,
+            )
+            self.servers_mock.get.assert_called_with(self.server.id)
+            self.images_mock.get.assert_called_with(self.image.id)
+            self.server.rebuild.assert_called_with(*args, **kwargs)
+
+    def test_rebuild_with_key_name_and_unset(self):
+        self.server.key_name = 'mykey'
+        arglist = [
+            self.server.id,
+            '--key-name', self.server.key_name,
+            '--key-unset',
+        ]
+        verifylist = [
+            ('server', self.server.id),
+            ('key_name', self.server.key_name)
+        ]
+        self.assertRaises(utils.ParserException,
+                          self.check_parser,
+                          self.cmd, arglist, verifylist)
+
 
 class TestServerRemoveFixedIP(TestServer):