From 1d78a8bbc5c8124b8fd739db9249a17af95661ef Mon Sep 17 00:00:00 2001
From: Andrew Phelps <andrew.phelps@canonical.com>
Date: Thu, 3 Oct 2024 16:49:20 -0400
Subject: [PATCH] internal/commands, internal/statemachine: allow forwarding
 components to "snap prepare-image"

---
 internal/commands/snap.go            |  1 +
 internal/statemachine/snap_states.go |  1 +
 internal/statemachine/snap_test.go   | 55 ++++++++++++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/internal/commands/snap.go b/internal/commands/snap.go
index 52109cb8..82b3b5fe 100644
--- a/internal/commands/snap.go
+++ b/internal/commands/snap.go
@@ -13,6 +13,7 @@ type SnapOpts struct {
 	AppArmorKernelFeaturesDir string         `long:"apparmor-features-dir" description:"Optional path to apparmor kernel features directory"`
 	PreseedSignKey            string         `long:"preseed-sign-key" description:"Name of the key to use to sign preseed assertion, otherwise use the default key"`
 	Snaps                     []string       `long:"snap" description:"Install extra snaps. These are passed through to \"snap prepare-image\". The snap argument can include additional information about the channel and/or risk with the following syntax: <snap>=<channel|risk>" value-name:"SNAP"`
+	Components                []string       `long:"comp" description:"Install extra components. These are passed through to \"snap prepare-image\"." value-name:"COMPONENT"`
 	CloudInit                 string         `long:"cloud-init" description:"cloud-config data to be copied to the image" value-name:"USER-DATA-FILE"`
 	Revisions                 map[string]int `long:"revision" description:"The revision of a specific snap to install in the image." value-name:"REVISION"`
 	SysfsOverlay              string         `long:"sysfs-overlay" description:"The optional sysfs overlay to used for preseeding. Directories from /sys/class/* and /sys/devices/platform will be bind-mounted to the chroot when preseeding"`
diff --git a/internal/statemachine/snap_states.go b/internal/statemachine/snap_states.go
index 5f4cb6ae..18f2ec61 100644
--- a/internal/statemachine/snap_states.go
+++ b/internal/statemachine/snap_states.go
@@ -28,6 +28,7 @@ func (stateMachine *StateMachine) prepareImage() error {
 		PrepareDir:                snapStateMachine.tempDirs.unpack,
 		Channel:                   snapStateMachine.commonFlags.Channel,
 		Customizations:            snapStateMachine.imageOptsCustomizations(),
+		Components:                snapStateMachine.Opts.Components,
 	}
 
 	var err error
diff --git a/internal/statemachine/snap_test.go b/internal/statemachine/snap_test.go
index 32d7c9b8..7c8c4407 100644
--- a/internal/statemachine/snap_test.go
+++ b/internal/statemachine/snap_test.go
@@ -235,6 +235,61 @@ func TestSuccessfulSnapCore20(t *testing.T) {
 	asserter.AssertErrNil(err, true)
 }
 
+func TestSuccessfulSnapCore20WithComponents(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode.")
+	}
+	asserter := helper.Asserter{T: t}
+	restoreCWD := testhelper.SaveCWD()
+	defer restoreCWD()
+
+	var stateMachine SnapStateMachine
+	stateMachine.commonFlags, stateMachine.stateMachineFlags = helper.InitCommonOpts()
+	stateMachine.parent = &stateMachine
+	// note that we must use a dangerous model here, since we're testing
+	// explicitly adding components
+	stateMachine.Args.ModelAssertion = filepath.Join("testdata", "modelAssertion20Dangerous")
+	stateMachine.Opts.FactoryImage = true
+	stateMachine.Opts.Snaps = []string{"core24", "test-snap-with-components"}
+	stateMachine.Opts.Components = []string{"test-snap-with-components+one", "test-snap-with-components+two"}
+	workDir, err := os.MkdirTemp("/tmp", "ubuntu-image-")
+	asserter.AssertErrNil(err, true)
+	t.Cleanup(func() { os.RemoveAll(workDir) })
+	stateMachine.stateMachineFlags.WorkDir = workDir
+
+	err = stateMachine.Setup()
+	asserter.AssertErrNil(err, true)
+
+	err = stateMachine.Run()
+	asserter.AssertErrNil(err, true)
+
+	for _, glob := range []string{"test-snap-with-components+one_*.comp", "test-snap-with-components+two_*.comp", "test-snap-with-components_*.snap"} {
+		matches, err := filepath.Glob(filepath.Join(
+			stateMachine.tempDirs.rootfs,
+			"systems/*/snaps/",
+			glob,
+		))
+		asserter.AssertErrNil(err, true)
+
+		if len(matches) != 1 {
+			t.Errorf("Expected exactly one match for %s, got %d", glob, len(matches))
+		}
+	}
+
+	// make sure the "factory" boot flag was set
+	grubenvFile := filepath.Join(stateMachine.tempDirs.rootfs,
+		"EFI", "ubuntu", "grubenv")
+	grubenvBytes, err := os.ReadFile(grubenvFile)
+	asserter.AssertErrNil(err, true)
+
+	if !strings.Contains(string(grubenvBytes), "snapd_boot_flags=factory") {
+		t.Errorf("grubenv file does not have factory boot flag set")
+	}
+
+	err = stateMachine.Teardown()
+	asserter.AssertErrNil(err, true)
+}
+
 // TestSuccessfulSnapCore18 builds a core 18 image with a few special options
 func TestSuccessfulSnapCore18(t *testing.T) {
 	if testing.Short() {