diff --git a/openstackclient/image/v1/image.py b/openstackclient/image/v1/image.py index 490b4122f6..91b0db4280 100644 --- a/openstackclient/image/v1/image.py +++ b/openstackclient/image/v1/image.py @@ -528,6 +528,16 @@ class SaveImage(command.Command): def get_parser(self, prog_name): parser = super().get_parser(prog_name) + parser.add_argument( + "--chunk-size", + type=int, + default=1024, + metavar="", + help=_( + "Size in bytes to read from the wire and buffer at one " + "time (default: 1024)" + ), + ) parser.add_argument( "--file", metavar="", @@ -550,7 +560,12 @@ class SaveImage(command.Command): if output_file is None: output_file = getattr(sys.stdout, "buffer", sys.stdout) - image_client.download_image(image.id, stream=True, output=output_file) + image_client.download_image( + image.id, + stream=True, + output=output_file, + chunk_size=parsed_args.chunk_size, + ) class SetImage(command.Command): diff --git a/openstackclient/image/v2/image.py b/openstackclient/image/v2/image.py index d2a1a74eae..9cd24f2ed1 100644 --- a/openstackclient/image/v2/image.py +++ b/openstackclient/image/v2/image.py @@ -1066,6 +1066,16 @@ class SaveImage(command.Command): def get_parser(self, prog_name): parser = super().get_parser(prog_name) + parser.add_argument( + "--chunk-size", + type=int, + default=1024, + metavar="", + help=_( + "Size in bytes to read from the wire and buffer at one " + "time (default: 1024)" + ), + ) parser.add_argument( "--file", metavar="", @@ -1090,7 +1100,12 @@ class SaveImage(command.Command): if output_file is None: output_file = getattr(sys.stdout, "buffer", sys.stdout) - image_client.download_image(image.id, stream=True, output=output_file) + image_client.download_image( + image.id, + stream=True, + output=output_file, + chunk_size=parsed_args.chunk_size, + ) class SetImage(command.Command): diff --git a/openstackclient/tests/unit/image/v1/test_image.py b/openstackclient/tests/unit/image/v1/test_image.py index 11d0830201..2d870ff89e 100644 --- a/openstackclient/tests/unit/image/v1/test_image.py +++ b/openstackclient/tests/unit/image/v1/test_image.py @@ -757,3 +757,50 @@ class TestImageShow(image_fakes.TestImagev1): size_index = columns.index('size') self.assertEqual(data[size_index].human_readable(), '2K') + + +class TestImageSave(image_fakes.TestImagev1): + image = image_fakes.create_one_image({}) + + def setUp(self): + super().setUp() + + self.image_client.find_image.return_value = self.image + self.image_client.download_image.return_value = self.image + + # Get the command object to test + self.cmd = image.SaveImage(self.app, None) + + def test_save_data(self): + arglist = ['--file', '/path/to/file', self.image.id] + + verifylist = [('file', '/path/to/file'), ('image', self.image.id)] + parsed_args = self.check_parser(self.cmd, arglist, verifylist) + + self.cmd.take_action(parsed_args) + + self.image_client.download_image.assert_called_once_with( + self.image.id, output='/path/to/file', stream=True, chunk_size=1024 + ) + + def test_save_data_with_chunk_size(self): + arglist = [ + '--file', + '/path/to/file', + '--chunk-size', + '2048', + self.image.id, + ] + + verifylist = [ + ('file', '/path/to/file'), + ('chunk_size', 2048), + ('image', self.image.id), + ] + parsed_args = self.check_parser(self.cmd, arglist, verifylist) + + self.cmd.take_action(parsed_args) + + self.image_client.download_image.assert_called_once_with( + self.image.id, output='/path/to/file', stream=True, chunk_size=2048 + ) diff --git a/openstackclient/tests/unit/image/v2/test_image.py b/openstackclient/tests/unit/image/v2/test_image.py index 4f5db9f002..2389d685fa 100644 --- a/openstackclient/tests/unit/image/v2/test_image.py +++ b/openstackclient/tests/unit/image/v2/test_image.py @@ -2228,7 +2228,29 @@ class TestImageSave(TestImage): self.cmd.take_action(parsed_args) self.image_client.download_image.assert_called_once_with( - self.image.id, stream=True, output='/path/to/file' + self.image.id, output='/path/to/file', stream=True, chunk_size=1024 + ) + + def test_save_data_with_chunk_size(self): + arglist = [ + '--file', + '/path/to/file', + '--chunk-size', + '2048', + self.image.id, + ] + + verifylist = [ + ('filename', '/path/to/file'), + ('chunk_size', 2048), + ('image', self.image.id), + ] + parsed_args = self.check_parser(self.cmd, arglist, verifylist) + + self.cmd.take_action(parsed_args) + + self.image_client.download_image.assert_called_once_with( + self.image.id, output='/path/to/file', stream=True, chunk_size=2048 ) diff --git a/releasenotes/notes/add-chunk-size-to-image-save-37871f9e62693264.yaml b/releasenotes/notes/add-chunk-size-to-image-save-37871f9e62693264.yaml new file mode 100644 index 0000000000..433dc87733 --- /dev/null +++ b/releasenotes/notes/add-chunk-size-to-image-save-37871f9e62693264.yaml @@ -0,0 +1,5 @@ +--- +features: + - | + Add ``--chunk-size`` option to ``image save`` command to control the size + of bytes to read at one time.