diff --git a/cmd/image-builder/cmd.go b/cmd/image-builder/cmd.go index fb37349d1d..8d703e5f9b 100644 --- a/cmd/image-builder/cmd.go +++ b/cmd/image-builder/cmd.go @@ -6,6 +6,7 @@ import ( "os" "github.com/osbuild/image-builder/internal/olog" + "github.com/osbuild/image-builder/pkg/datasizes" ilog "github.com/osbuild/image-builder/pkg/olog" "github.com/spf13/cobra" "github.com/spf13/cobra/doc" @@ -191,6 +192,8 @@ func setupManifestCmd() (*cobra.Command, error) { manifestCmd := &cobra.Command{ Use: "manifest ", Short: "Build manifest for the given image-type, e.g. qcow2 (tip: combine with --distro, --arch)", + Long: "Build a manifest for the selected image type. The --image-size flag accepts bytes or a value with a data-size unit.", + Example: " image-builder manifest qcow2 --image-size \"1 GiB\"", RunE: cmdManifest, SilenceUsage: true, Args: cobra.ExactArgs(1), @@ -209,7 +212,8 @@ func setupManifestCmd() (*cobra.Command, error) { manifestCmd.Flags().String("bootc-default-fs", "", `default filesystem to use for the bootc install (e.g. ext4)`) manifestCmd.Flags().Bool("bootc-no-default-kernel-args", false, `don't use the default kernel arguments`) manifestCmd.Flags().Bool("bootc-pull-container", false, `pull bootc container from remote location instead of using it from local container storage`) - manifestCmd.Flags().Uint64("image-size", 0, `override the default image size in bytes`) + var imageSize datasizes.Size + manifestCmd.Flags().TextVar(&imageSize, "image-size", imageSize, `override the default image size (e.g. 1 GiB)`) manifestCmd.Flags().Bool("use-librepo", true, `use librepo to download packages (disable if you use old versions of osbuild)`) if err := manifestCmd.Flags().MarkHidden("use-librepo"); err != nil { return nil, err diff --git a/cmd/image-builder/cmd_internal_test.go b/cmd/image-builder/cmd_internal_test.go new file mode 100644 index 0000000000..d7ad390e64 --- /dev/null +++ b/cmd/image-builder/cmd_internal_test.go @@ -0,0 +1,49 @@ +package main + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/osbuild/image-builder/pkg/datasizes" +) + +func TestManifestImageSizeFlag(t *testing.T) { + cases := []struct { + name string + input string + expected datasizes.Size + }{ + { + name: "bytes", + input: "1073741824", + expected: datasizes.Size(datasizes.GiB), + }, + { + name: "with-unit", + input: "1 GiB", + expected: datasizes.Size(datasizes.GiB), + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + manifestCmd, err := setupManifestCmd() + require.NoError(t, err) + require.NoError(t, manifestCmd.ParseFlags([]string{"--image-size", tc.input})) + + var imageSize datasizes.Size + require.NoError(t, manifestCmd.Flags().GetText("image-size", &imageSize)) + assert.Equal(t, tc.expected, imageSize) + }) + } +} + +func TestManifestCommandDocumentsImageSizeUsage(t *testing.T) { + manifestCmd, err := setupManifestCmd() + require.NoError(t, err) + + assert.Contains(t, manifestCmd.Long, "--image-size") + assert.Contains(t, manifestCmd.Example, `--image-size "1 GiB"`) +} diff --git a/cmd/image-builder/main.go b/cmd/image-builder/main.go index 2d3fac6db0..a3da3dcbc8 100644 --- a/cmd/image-builder/main.go +++ b/cmd/image-builder/main.go @@ -21,6 +21,7 @@ import ( "github.com/osbuild/image-builder/pkg/bootc" "github.com/osbuild/image-builder/pkg/cloud" "github.com/osbuild/image-builder/pkg/customizations/subscription" + "github.com/osbuild/image-builder/pkg/datasizes" "github.com/osbuild/image-builder/pkg/distro" "github.com/osbuild/image-builder/pkg/distro/generic" "github.com/osbuild/image-builder/pkg/imagefilter" @@ -435,7 +436,8 @@ func generateManifest(pbar progress.ProgressBar, cmd *cobra.Command, args []stri if err != nil { return nil, err } - imageSize, err := cmd.Flags().GetUint64("image-size") + var imageSize datasizes.Size + err = cmd.Flags().GetText("image-size", &imageSize) if err != nil { return nil, err } @@ -567,7 +569,7 @@ func generateManifest(pbar progress.ProgressBar, cmd *cobra.Command, args []stri Facts: &facts.ImageOptions{APIType: facts.IBCLI_APITYPE}, OSTree: ostreeImgOpts, Subscription: subscription, - Size: imageSize, + Size: imageSize.Uint64(), Bootc: &distro.BootcImageOptions{ InstallerPayloadRef: bootcInstallerPayloadRef, OmitDefaultKernelArgs: bootcOmitDefaultKernelArgs, diff --git a/pkg/datasizes/size.go b/pkg/datasizes/size.go index cde6479049..c474b3f0e7 100644 --- a/pkg/datasizes/size.go +++ b/pkg/datasizes/size.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/json" "fmt" + "strconv" "github.com/osbuild/image-builder/internal/common" ) @@ -19,6 +20,19 @@ func (si Size) Uint64() uint64 { return uint64(si) } +func (si *Size) UnmarshalText(data []byte) error { + value, err := Parse(string(data)) + if err != nil { + return fmt.Errorf("error decoding size: %w", err) + } + *si = Size(value) + return nil +} + +func (si Size) MarshalText() ([]byte, error) { + return []byte(strconv.FormatUint(si.Uint64(), 10)), nil +} + func (si *Size) UnmarshalTOML(data interface{}) error { i, err := decodeSize(data) if err != nil { diff --git a/pkg/datasizes/size_test.go b/pkg/datasizes/size_test.go index 1830a73009..361dcb2aa4 100644 --- a/pkg/datasizes/size_test.go +++ b/pkg/datasizes/size_test.go @@ -166,3 +166,42 @@ func TestSizeUnmarshalHappy(t *testing.T) { func TestSizeUint64(t *testing.T) { assert.Equal(t, datasizes.Size(1234).Uint64(), uint64(1234)) } + +func TestSizeTextMarshaling(t *testing.T) { + cases := []struct { + name string + input string + expected datasizes.Size + }{ + { + name: "bytes", + input: "1073741824", + expected: datasizes.Size(datasizes.GiB), + }, + { + name: "with-unit", + input: "1 GiB", + expected: datasizes.Size(datasizes.GiB), + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + var size datasizes.Size + assert.NoError(t, size.UnmarshalText([]byte(tc.input))) + assert.Equal(t, tc.expected, size) + + text, err := size.MarshalText() + assert.NoError(t, err) + assert.Equal(t, "1073741824", string(text)) + }) + } +} + +func TestSizeUnmarshalTextUnhappy(t *testing.T) { + var size datasizes.Size + + err := size.UnmarshalText([]byte("20 KG")) + + assert.EqualError(t, err, "error decoding size: unknown data size units in string: 20 KG") +}