manifest: accept human-readable image sizes

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 <image-type>",
 		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

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"`)
+}

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,

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 {

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")
+}