o/devicestate: introduce additional mocking functions for gadget functions

overlord/devicestate/export_test.go

@@ -21,6 +21,7 @@ package devicestate
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"os/user"
 	"time"
@@ -366,6 +367,38 @@ func MockInstallFactoryReset(f func(model gadget.Model, gadgetRoot, kernelRoot,
 	return restore
 }
 
+func MockInstallWriteContent(f func(onVolumes map[string]*gadget.Volume, allLaidOutVols map[string]*gadget.LaidOutVolume, encSetupData *install.EncryptionSetupData, observer gadget.ContentObserver, perfTimings timings.Measurer) ([]*gadget.OnDiskVolume, error)) (restore func()) {
+	old := installWriteContent
+	installWriteContent = f
+	return func() {
+		installWriteContent = old
+	}
+}
+
+func MockInstallMountVolumes(f func(onVolumes map[string]*gadget.Volume, encSetupData *install.EncryptionSetupData) (espMntDir string, unmount func() error, err error)) (restore func()) {
+	old := installMountVolumes
+	installMountVolumes = f
+	return func() {
+		installMountVolumes = old
+	}
+}
+
+func MockInstallEncryptPartitions(f func(onVolumes map[string]*gadget.Volume, encryptionType secboot.EncryptionType, model *asserts.Model, gadgetRoot, kernelRoot string, perfTimings timings.Measurer) (*install.EncryptionSetupData, error)) (restore func()) {
+	old := installEncryptPartitions
+	installEncryptPartitions = f
+	return func() {
+		installEncryptPartitions = old
+	}
+}
+
+func MockInstallSaveStorageTraits(f func(model gadget.Model, allLaidOutVols map[string]*gadget.LaidOutVolume, encryptSetupData *install.EncryptionSetupData) error) (restore func()) {
+	old := installSaveStorageTraits
+	installSaveStorageTraits = f
+	return func() {
+		installSaveStorageTraits = old
+	}
+}
+
 func MockSecbootStageEncryptionKeyChange(f func(node string, key keys.EncryptionKey) error) (restore func()) {
 	restore = testutil.Backup(&secbootStageEncryptionKeyChange)
 	secbootStageEncryptionKeyChange = f
@@ -491,3 +524,26 @@ func MockCreateAllKnownSystemUsers(createAllUsers func(state *state.State, asser
 	createAllKnownSystemUsers = createAllUsers
 	return restore
 }
+
+func MockEncryptionSetupDataInCache(st *state.State, label string) (restore func()) {
+	encSetup := &install.EncryptionSetupData{}
+	key := encryptionSetupDataKey{label}
+	st.Cache(key, encSetup)
+	return func() { st.Cache(key, nil) }
+}
+
+func CheckEncryptionSetupDataFromCache(st *state.State, label string) error {
+	cached := st.Cached(encryptionSetupDataKey{label})
+	if cached == nil {
+		return fmt.Errorf("no EncryptionSetupData found in cache")
+	}
+	if _, ok := cached.(*install.EncryptionSetupData); !ok {
+		return fmt.Errorf("wrong data type under encryptionSetupDataKey")
+	}
+	return nil
+}
+
+func CleanUpEncryptionSetupDataInCache(st *state.State, label string) {
+	key := encryptionSetupDataKey{label}
+	st.Cache(key, nil)
+}