add tests

Author: himaja-kesari

toolkit/tools/imagecreator/main.go

@@ -18,11 +18,12 @@ import (
 )
 
 type ImageCreatorCmd struct {
-	BuildDir        string   `name:"build-dir" help:"Directory to run build out of." required:""`
-	ConfigFile      string   `name:"config-file" help:"Path of the image customization config file." required:""`
-	RpmSources      []string `name:"rpm-source" help:"Path to a RPM repo config file or a directory containing RPMs."`
-	ToolsTar        string   `name:"tools-file" help:"Path to tdnf worker tarball"`
-	OutputImageFile string   `name:"output-image-file" help:"Path to write the customized image to."`
+	BuildDir          string   `name:"build-dir" help:"Directory to run build out of." required:""`
+	ConfigFile        string   `name:"config-file" help:"Path of the image customization config file." required:""`
+	RpmSources        []string `name:"rpm-source" help:"Path to a RPM repo config file or a directory containing RPMs."`
+	ToolsTar          string   `name:"tools-file" help:"Path to tdnf worker tarball"`
+	OutputImageFile   string   `name:"output-image-file" help:"Path to write the customized image to."`
+	OutputImageFormat string   `name:"output-image-format" placeholder:"(vhd|vhd-fixed|vhdx|qcow2|raw)" help:"Format of output image." enum:"${imageformat}" default:""`
 	exekong.LogFlags
 }
 
@@ -44,7 +45,7 @@ func main() {
 
 	logger.InitBestEffort(ptrutils.PtrTo(cli.LogFlags.AsLoggerFlags()))
 
-	err := imagecreatorlib.CreateImageWithConfigFile(cli.BuildDir, cli.ConfigFile, cli.RpmSources, cli.ToolsTar, cli.OutputImageFile)
+	err := imagecreatorlib.CreateImageWithConfigFile(cli.BuildDir, cli.ConfigFile, cli.RpmSources, cli.ToolsTar, cli.OutputImageFile, cli.OutputImageFormat)
 	if err != nil {
 		log.Fatalf("image creation failed:\n%v", err)
 	}

toolkit/tools/pkg/imagecustomizerlib/testdata/testrpms/download-test-rpms.sh renamed to toolkit/tools/internal/testutils/testrpms/download-test-rpms.sh

@@ -6,28 +6,39 @@ CONTAINER_TAG="imagecustomizertestrpms:latest"
 DOCKERFILE_DIR="$SCRIPT_DIR/downloader"
 
 AZURELINUX_2_CONTAINER_IMAGE="mcr.microsoft.com/cbl-mariner/base/core:2.0"
+AZURELINUX_3_CONTAINER_IMAGE="mcr.microsoft.com/azurelinux/base/core:3.0"
 
 IMAGE_VERSION="2.0"
 
-while getopts "t:" flag
+IMAGE_CREATOR="false"
+TOOLS_FILE="$SCRIPT_DIR/tools.tar.gz"
+
+
+while getopts "s:t:" flag
 do
     case "${flag}" in
+        s) IMAGE_CREATOR="$OPTARG";;
         t) IMAGE_VERSION="$OPTARG";;
         h) ;;&
         ?)
-            echo "Usage: download-test-rpms.sh [-t IMAGE_VERSION]"
+            echo "Usage: download-test-rpms.sh [-t IMAGE_VERSION] [-s IMAGE_CREATOR]"
             echo ""
             echo "Args:"
             echo "  -t IMAGE_VERSION   The Azure Image version to download the RPMs for."
+            echo "  -s IMAGE_CREATOR    If set to true, the script will create a tar.gz file with the tools and download the rpms needed to test imagecreator."
             echo "  -h Show help"
             exit 1;;
     esac
 done
 
+
 case "${IMAGE_VERSION}" in
+  3.0)
+    CONTAINER_IMAGE="$AZURELINUX_3_CONTAINER_IMAGE"
+    ;;
   2.0)
     CONTAINER_IMAGE="$AZURELINUX_2_CONTAINER_IMAGE"
-    ;;
+    ;;  
   *)
     echo "error: unsupported Azure Linux version: $IMAGE_VERSION"
     exit 1;;
@@ -44,10 +55,19 @@ mkdir -p "$OUT_DIR"
 
 # Build a container image that contains the RPMs.
 docker build \
-  --build-arg "baseimage=$AZURELINUX_2_CONTAINER_IMAGE" \
+  --build-arg "baseimage=$CONTAINER_IMAGE" \
+  --build-arg "imagecreator=$IMAGE_CREATOR" \
   --tag "$CONTAINER_TAG" \
   "$DOCKERFILE_DIR"
 
+
+# if IMAGE_CREATOR is set to true, create a tar.gz file with the tools.
+if [ "$IMAGE_CREATOR" = "true" ]; then
+  docker create --name temp-container ${CONTAINER_TAG}
+  docker export temp-container | gzip > "$TOOLS_FILE"
+  docker rm temp-container
+fi
+
 # Extract the RPM files.
 docker run \
   --rm \

toolkit/tools/internal/testutils/testrpms/downloader/Dockerfile

@@ -0,0 +1,27 @@
+ARG baseimage
+FROM ${baseimage}
+
+ARG imagecreator
+
+
+
+RUN tdnf update -y && \
+   tdnf install -y dnf dnf-plugins-core createrepo_c
+
+RUN dnf download -y --resolve --alldeps --destdir /downloadedrpms \
+   jq golang 
+   
+RUN if [ "${imagecreator}" = "true" ]; then \
+
+    dnf download -y --resolve --alldeps --destdir /downloadedrpms \
+    azurelinux-release \
+    azurelinux-repos azurelinux-rpm-macros bash dbus dracut-hostonly e2fsprogs filesystem \
+    grub2 grub2-efi-binary iana-etc iproute iputils irqbalance \
+    ncurses-libs openssl rpm rpm-libs shadow-utils shim sudo \
+    systemd systemd-networkd systemd-resolved systemd-udev tdnf \
+    tdnf-plugin-repogpgcheck util-linux zlib kernel initramfs ; \
+
+    fi
+
+# Add repo metadata, so that the directory can be used in a .repo file.
+RUN createrepo --compatibility --update /downloadedrpms

toolkit/tools/internal/testutils/testutils.go

@@ -0,0 +1,115 @@
+package testutils
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/microsoft/azurelinux/toolkit/tools/internal/file"
+	"github.com/stretchr/testify/assert"
+)
+
+func GetImageFileType(filePath string) (string, error) {
+	file, err := os.OpenFile(filePath, os.O_RDONLY, 0)
+	if err != nil {
+		return "", err
+	}
+	defer file.Close()
+
+	stat, err := file.Stat()
+	if err != nil {
+		return "", err
+	}
+
+	firstBytes := make([]byte, 512)
+	firstBytesCount, err := file.Read(firstBytes)
+	if err != nil {
+		return "", err
+	}
+
+	lastBytes := make([]byte, 512)
+	lastBytesCount, err := file.ReadAt(lastBytes, max(0, stat.Size()-512))
+	if err != nil {
+		return "", err
+	}
+
+	switch {
+	case firstBytesCount >= 8 && bytes.Equal(firstBytes[:8], []byte("conectix")):
+		return "vhd", nil
+
+	case firstBytesCount >= 8 && bytes.Equal(firstBytes[:8], []byte("vhdxfile")):
+		return "vhdx", nil
+
+	case isZstFile(firstBytes):
+		return "zst", nil
+
+	// Check for the MBR signature (which exists even on GPT formatted drives).
+	case firstBytesCount >= 512 && bytes.Equal(firstBytes[510:512], []byte{0x55, 0xAA}):
+		switch {
+		case lastBytesCount >= 512 && bytes.Equal(lastBytes[:8], []byte("conectix")):
+			return "vhd-fixed", nil
+
+		default:
+			return "raw", nil
+		}
+
+	default:
+		return "", fmt.Errorf("unknown file type: %s", filePath)
+	}
+}
+
+func isZstFile(firstBytes []byte) bool {
+	if len(firstBytes) < 4 {
+		return false
+	}
+
+	magicNumber := binary.LittleEndian.Uint32(firstBytes[:4])
+
+	// 0xFD2FB528 is a zst frame.
+	// 0x184D2A50-0x184D2A5F are skippable ztd frames.
+	return magicNumber == 0xFD2FB528 || (magicNumber >= 0x184D2A50 && magicNumber <= 0x184D2A5F)
+}
+
+func GetDownloadedRpmsDir(t *testing.T, testutilsDir string, azureLinuxVersion string, imagecreator bool) string {
+	downloadedRpmsDir := filepath.Join(testutilsDir, "testrpms/downloadedrpms", azureLinuxVersion)
+	dirExists, err := file.DirExists(downloadedRpmsDir)
+	if !assert.NoErrorf(t, err, "cannot access downloaded RPMs dir (%s)", downloadedRpmsDir) {
+		t.FailNow()
+	}
+	if !assert.True(t, dirExists) {
+		// log the downloadedRpmsDir
+		t.Logf("downloadedRpmsDir: %s", downloadedRpmsDir)
+		t.Logf("test requires offline RPMs")
+		t.Logf("please run toolkit/tools/internal/testutils/testrpms/download-test-rpms.sh -t %s -s %t",
+			azureLinuxVersion, imagecreator)
+		t.FailNow()
+	}
+
+	return downloadedRpmsDir
+}
+
+func GetDownloadedToolsFile(t *testing.T, testutilsDir string, azureLinuxVersion string, imagecreator bool) string {
+	GetDownloadedToolsFile := filepath.Join(testutilsDir, "testrpms/tools.tar.gz")
+	if !assert.FileExists(t, GetDownloadedToolsFile) {
+		t.Logf("test requires downloaded tools file")
+		t.Logf("please run toolkit/tools/internal/testutils/testrpms/download-test-rpms.sh -t %s -s %t",
+			azureLinuxVersion, imagecreator)
+		t.FailNow()
+	}
+	return GetDownloadedToolsFile
+}
+
+func GetDownloadedRpmsRepoFile(t *testing.T, testutilsDir string, azureLinuxVersion string, withGpgKey bool, imagecreator bool) string {
+	dir := GetDownloadedRpmsDir(t, testutilsDir, azureLinuxVersion, imagecreator)
+
+	suffix := "nokey"
+	if withGpgKey {
+		suffix = "withkey"
+	}
+
+	repoFile := filepath.Join(dir, "../", fmt.Sprintf("rpms-%s-%s.repo", azureLinuxVersion, suffix))
+	return repoFile
+}