Feedback updates

Author: cwize1

toolkit/tools/pkg/imagecustomizerlib/distrohandler_test.go

@@ -161,7 +161,7 @@ func testCustomizeImageUnsupportedPackageSnapshotTimeHelper(t *testing.T, baseIm
 }
 
 func TestCustomizeImageUnsupportedRpmSources(t *testing.T) {
-	for _, baseImageInfo := range slices.Concat(baseImageUbuntuAll) {
+	for _, baseImageInfo := range baseImageUbuntuAll {
 		t.Run(baseImageInfo.Name, func(t *testing.T) {
 			testCustomizeImageUnsupportedRpmSourcesHelper(t, baseImageInfo)
 		})

toolkit/tools/pkg/imagecustomizerlib/distrohandler_ubuntu.go

@@ -31,6 +31,11 @@ const (
 	grubEfiPackageDebianArm64 = "grub-efi-arm64"
 )
 
+var (
+	ErrUbuntuUnsupportedBootloaderHardReset   = NewImageCustomizerError("Validation:UbuntuUnsupportedBootloaderHardReset", "bootloader hard-reset API is not supported yet for Ubuntu images")
+	ErrUbuntuUnsupportedDisableBaseImageRepos = NewImageCustomizerError("Validation:UbuntuUnsupportedDisableBaseImageRepos", "disabling base image package repositories is not supported yet for Ubuntu images")
+)
+
 func newUbuntuDistroHandler(targetOs targetos.TargetOs) *ubuntuDistroHandler {
 	logger.Log.Debugf("Distro handler: Ubuntu (distro='%s', versionid='%s')", targetOs.Distro, targetOs.VersionId)
 
@@ -72,7 +77,7 @@ func (d *ubuntuDistroHandler) checkForUnsupportedApis(rc *ResolvedConfig) error
 	// Check if Ubuntu is being used with bootloader hard-reset.
 	// Ubuntu bootloader config logic is not yet fully implemented.
 	if rc.BootLoader.ResetType == imagecustomizerapi.ResetBootLoaderTypeHard {
-		return fmt.Errorf("bootloader hard-reset API is not yet supported for Ubuntu images")
+		return ErrUbuntuUnsupportedBootloaderHardReset
 	}
 
 	if len(rc.Options.RpmsSources) > 0 {
@@ -83,7 +88,7 @@ func (d *ubuntuDistroHandler) checkForUnsupportedApis(rc *ResolvedConfig) error
 	// passes --disable-base-image-rpm-repos. Ubuntu does not use RPM repos, so disabling
 	// them is not meaningful and likely indicates a configuration mistake.
 	if !rc.Options.UseBaseImageRpmRepos {
-		return fmt.Errorf("disabling base image package repositories is not supported yet for Ubuntu images")
+		return ErrUbuntuUnsupportedDisableBaseImageRepos
 	}
 
 	if rc.HasPackageSnapshotTime() {
@@ -202,7 +207,7 @@ func (d *ubuntuDistroHandler) ConfigureDiskBootLoader(imageConnection *imageconn
 	selinuxConfig imagecustomizerapi.SELinux, kernelCommandLine imagecustomizerapi.KernelCommandLine,
 	currentSELinuxMode imagecustomizerapi.SELinuxMode, newImage bool,
 ) error {
-	return fmt.Errorf("bootloader hard-reset API is not yet supported for Ubuntu images")
+	return ErrUbuntuUnsupportedBootloaderHardReset
 }
 
 func (d *ubuntuDistroHandler) ReadGrubConfigLinuxArgs(bootDir string) (map[string][]grubConfigLinuxArg, error) {

toolkit/tools/pkg/imagecustomizerlib/distrohandler_ubuntu_test.go

@@ -0,0 +1,56 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+package imagecustomizerlib
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/azure-linux-image-tools/toolkit/tools/imagecustomizerapi"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCustomizeImageUbuntuUnsupportedAPIs(t *testing.T) {
+	for _, baseImageInfo := range baseImageUbuntuAll {
+		t.Run(baseImageInfo.Name, func(t *testing.T) {
+			testCustomizeImageUbuntuUnsupportedAPIsHelper(t, baseImageInfo)
+		})
+	}
+}
+
+func testCustomizeImageUbuntuUnsupportedAPIsHelper(t *testing.T, baseImageInfo testBaseImageInfo) {
+	baseImage := checkSkipForCustomizeImage(t, baseImageInfo)
+
+	testTmpDir := filepath.Join(tmpDir, "TestCustomizeImageUbuntuUnsupportedAPIs_"+baseImageInfo.Name)
+	defer os.RemoveAll(testTmpDir)
+
+	buildDir := filepath.Join(testTmpDir, "build")
+
+	options := ImageCustomizerOptions{
+		BuildDir:             buildDir,
+		InputImageFile:       baseImage,
+		OutputImageFile:      "./out/image.vhdx",
+		OutputImageFormat:    "vhdx",
+		UseBaseImageRpmRepos: true,
+		PreviewFeatures:      baseImageInfo.PreviewFeatures,
+	}
+
+	config := &imagecustomizerapi.Config{
+		OS: &imagecustomizerapi.OS{},
+	}
+
+	config.OS.BootLoader.ResetType = imagecustomizerapi.ResetBootLoaderTypeHard
+
+	err := CustomizeImage(t.Context(), testTmpDir, config, options)
+	assert.ErrorIs(t, err, ErrUbuntuUnsupportedBootloaderHardReset)
+	assert.ErrorIs(t, err, ErrUnsupportedDistroApi)
+
+	config.OS.BootLoader.ResetType = imagecustomizerapi.ResetBootLoaderTypeDefault
+	options.UseBaseImageRpmRepos = false
+
+	err = CustomizeImage(t.Context(), testTmpDir, config, options)
+	assert.ErrorIs(t, err, ErrUbuntuUnsupportedDisableBaseImageRepos)
+	assert.ErrorIs(t, err, ErrUnsupportedDistroApi)
+}