gopkg.in/ubuntu-core/snappy.v0@v0.0.0-20210902073436-25a8614f10a6/overlord/devicestate/export_test.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2016-2018 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package devicestate
    21  
    22  import (
    23  	"context"
    24  	"net/http"
    25  	"time"
    26  
    27  	"github.com/snapcore/snapd/asserts"
    28  	"github.com/snapcore/snapd/boot"
    29  	"github.com/snapcore/snapd/gadget"
    30  	"github.com/snapcore/snapd/gadget/install"
    31  	"github.com/snapcore/snapd/httputil"
    32  	"github.com/snapcore/snapd/kernel/fde"
    33  	"github.com/snapcore/snapd/overlord/snapstate"
    34  	"github.com/snapcore/snapd/overlord/state"
    35  	"github.com/snapcore/snapd/overlord/storecontext"
    36  	"github.com/snapcore/snapd/sysconfig"
    37  	"github.com/snapcore/snapd/timings"
    38  )
    39  
    40  func MockKeyLength(n int) (restore func()) {
    41  	if n < 1024 {
    42  		panic("key length must be >= 1024")
    43  	}
    44  
    45  	oldKeyLength := keyLength
    46  	keyLength = n
    47  	return func() {
    48  		keyLength = oldKeyLength
    49  	}
    50  }
    51  
    52  func MockBaseStoreURL(url string) (restore func()) {
    53  	oldURL := baseStoreURL
    54  	baseStoreURL = mustParse(url).ResolveReference(authRef)
    55  	return func() {
    56  		baseStoreURL = oldURL
    57  	}
    58  }
    59  
    60  func MockRetryInterval(interval time.Duration) (restore func()) {
    61  	old := retryInterval
    62  	retryInterval = interval
    63  	return func() {
    64  		retryInterval = old
    65  	}
    66  }
    67  
    68  func MockMaxTentatives(max int) (restore func()) {
    69  	old := maxTentatives
    70  	maxTentatives = max
    71  	return func() {
    72  		maxTentatives = old
    73  	}
    74  }
    75  
    76  func MockTimeNow(f func() time.Time) (restore func()) {
    77  	old := timeNow
    78  	timeNow = f
    79  	return func() {
    80  		timeNow = old
    81  	}
    82  }
    83  
    84  func KeypairManager(m *DeviceManager) (keypairMgr asserts.KeypairManager) {
    85  	// XXX expose the with... method at some point
    86  	err := m.withKeypairMgr(func(km asserts.KeypairManager) error {
    87  		keypairMgr = km
    88  		return nil
    89  	})
    90  	if err != nil {
    91  		panic(err)
    92  	}
    93  	return keypairMgr
    94  }
    95  
    96  func SaveAvailable(m *DeviceManager) bool {
    97  	return m.saveAvailable
    98  }
    99  
   100  func SetSaveAvailable(m *DeviceManager, avail bool) {
   101  	m.saveAvailable = avail
   102  }
   103  
   104  func EnsureOperationalShouldBackoff(m *DeviceManager, now time.Time) bool {
   105  	return m.ensureOperationalShouldBackoff(now)
   106  }
   107  
   108  func BecomeOperationalBackoff(m *DeviceManager) time.Duration {
   109  	return m.becomeOperationalBackoff
   110  }
   111  
   112  func SetLastBecomeOperationalAttempt(m *DeviceManager, t time.Time) {
   113  	m.lastBecomeOperationalAttempt = t
   114  }
   115  
   116  func SetSystemMode(m *DeviceManager, mode string) {
   117  	m.sysMode = mode
   118  }
   119  
   120  func SetTimeOnce(m *DeviceManager, name string, t time.Time) error {
   121  	return m.setTimeOnce(name, t)
   122  }
   123  
   124  func PreloadGadget(m *DeviceManager) (sysconfig.Device, *gadget.Info, error) {
   125  	return m.preloadGadget()
   126  }
   127  
   128  func MockRepeatRequestSerial(label string) (restore func()) {
   129  	old := repeatRequestSerial
   130  	repeatRequestSerial = label
   131  	return func() {
   132  		repeatRequestSerial = old
   133  	}
   134  }
   135  
   136  func MockSnapstateInstallWithDeviceContext(f func(ctx context.Context, st *state.State, name string, opts *snapstate.RevisionOptions, userID int, flags snapstate.Flags, deviceCtx snapstate.DeviceContext, fromChange string) (*state.TaskSet, error)) (restore func()) {
   137  	old := snapstateInstallWithDeviceContext
   138  	snapstateInstallWithDeviceContext = f
   139  	return func() {
   140  		snapstateInstallWithDeviceContext = old
   141  	}
   142  }
   143  
   144  func MockSnapstateUpdateWithDeviceContext(f func(st *state.State, name string, opts *snapstate.RevisionOptions, userID int, flags snapstate.Flags, deviceCtx snapstate.DeviceContext, fromChange string) (*state.TaskSet, error)) (restore func()) {
   145  	old := snapstateUpdateWithDeviceContext
   146  	snapstateUpdateWithDeviceContext = f
   147  	return func() {
   148  		snapstateUpdateWithDeviceContext = old
   149  	}
   150  }
   151  
   152  func EnsureSeeded(m *DeviceManager) error {
   153  	return m.ensureSeeded()
   154  }
   155  
   156  func EnsureCloudInitRestricted(m *DeviceManager) error {
   157  	return m.ensureCloudInitRestricted()
   158  }
   159  
   160  var PopulateStateFromSeedImpl = populateStateFromSeedImpl
   161  
   162  type PopulateStateFromSeedOptions = populateStateFromSeedOptions
   163  
   164  func MockPopulateStateFromSeed(f func(*state.State, *PopulateStateFromSeedOptions, timings.Measurer) ([]*state.TaskSet, error)) (restore func()) {
   165  	old := populateStateFromSeed
   166  	populateStateFromSeed = f
   167  	return func() {
   168  		populateStateFromSeed = old
   169  	}
   170  }
   171  
   172  func EnsureBootOk(m *DeviceManager) error {
   173  	return m.ensureBootOk()
   174  }
   175  
   176  func SetBootOkRan(m *DeviceManager, b bool) {
   177  	m.bootOkRan = b
   178  }
   179  
   180  func SetInstalledRan(m *DeviceManager, b bool) {
   181  	m.ensureInstalledRan = b
   182  }
   183  
   184  func SetTriedSystemsRan(m *DeviceManager, b bool) {
   185  	m.ensureTriedRecoverySystemRan = b
   186  }
   187  
   188  func StartTime() time.Time {
   189  	return startTime
   190  }
   191  
   192  type (
   193  	RegistrationContext = registrationContext
   194  	RemodelContext      = remodelContext
   195  	SeededSystem        = seededSystem
   196  )
   197  
   198  func RegistrationCtx(m *DeviceManager, t *state.Task) (registrationContext, error) {
   199  	return m.registrationCtx(t)
   200  }
   201  
   202  func RemodelDeviceBackend(remodCtx remodelContext) storecontext.DeviceBackend {
   203  	return remodCtx.(interface {
   204  		deviceBackend() storecontext.DeviceBackend
   205  	}).deviceBackend()
   206  }
   207  
   208  func RemodelSetRecoverySystemLabel(remodCtx remodelContext, label string) {
   209  	remodCtx.setRecoverySystemLabel(label)
   210  }
   211  
   212  func RecordSeededSystem(m *DeviceManager, st *state.State, sys *seededSystem) error {
   213  	return m.recordSeededSystem(st, sys)
   214  }
   215  
   216  var (
   217  	LoadDeviceSeed               = loadDeviceSeed
   218  	UnloadDeviceSeed             = unloadDeviceSeed
   219  	ImportAssertionsFromSeed     = importAssertionsFromSeed
   220  	CheckGadgetOrKernel          = checkGadgetOrKernel
   221  	CheckGadgetValid             = checkGadgetValid
   222  	CheckGadgetRemodelCompatible = checkGadgetRemodelCompatible
   223  	CanAutoRefresh               = canAutoRefresh
   224  	NewEnoughProxy               = newEnoughProxy
   225  
   226  	IncEnsureOperationalAttempts = incEnsureOperationalAttempts
   227  	EnsureOperationalAttempts    = ensureOperationalAttempts
   228  
   229  	RemodelTasks = remodelTasks
   230  
   231  	RemodelCtx        = remodelCtx
   232  	CleanupRemodelCtx = cleanupRemodelCtx
   233  	CachedRemodelCtx  = cachedRemodelCtx
   234  
   235  	GadgetUpdateBlocked = gadgetUpdateBlocked
   236  	CurrentGadgetInfo   = currentGadgetInfo
   237  	PendingGadgetInfo   = pendingGadgetInfo
   238  
   239  	CriticalTaskEdges = criticalTaskEdges
   240  
   241  	CreateSystemForModelFromValidatedSnaps = createSystemForModelFromValidatedSnaps
   242  	LogNewSystemSnapFile                   = logNewSystemSnapFile
   243  	PurgeNewSystemSnapFiles                = purgeNewSystemSnapFiles
   244  	CreateRecoverySystemTasks              = createRecoverySystemTasks
   245  )
   246  
   247  func MockGadgetUpdate(mock func(current, update gadget.GadgetData, path string, policy gadget.UpdatePolicyFunc, observer gadget.ContentUpdateObserver) error) (restore func()) {
   248  	old := gadgetUpdate
   249  	gadgetUpdate = mock
   250  	return func() {
   251  		gadgetUpdate = old
   252  	}
   253  }
   254  
   255  func MockGadgetIsCompatible(mock func(current, update *gadget.Info) error) (restore func()) {
   256  	old := gadgetIsCompatible
   257  	gadgetIsCompatible = mock
   258  	return func() {
   259  		gadgetIsCompatible = old
   260  	}
   261  }
   262  
   263  func MockBootMakeSystemRunnable(f func(model *asserts.Model, bootWith *boot.BootableSet, seal *boot.TrustedAssetsInstallObserver) error) (restore func()) {
   264  	old := bootMakeRunnable
   265  	bootMakeRunnable = f
   266  	return func() {
   267  		bootMakeRunnable = old
   268  	}
   269  }
   270  
   271  func MockBootEnsureNextBootToRunMode(f func(systemLabel string) error) (restore func()) {
   272  	old := bootEnsureNextBootToRunMode
   273  	bootEnsureNextBootToRunMode = f
   274  	return func() {
   275  		bootEnsureNextBootToRunMode = old
   276  	}
   277  }
   278  
   279  func MockSecbootCheckTPMKeySealingSupported(f func() error) (restore func()) {
   280  	old := secbootCheckTPMKeySealingSupported
   281  	secbootCheckTPMKeySealingSupported = f
   282  	return func() {
   283  		secbootCheckTPMKeySealingSupported = old
   284  	}
   285  }
   286  
   287  func MockHttputilNewHTTPClient(f func(opts *httputil.ClientOptions) *http.Client) (restore func()) {
   288  	old := httputilNewHTTPClient
   289  	httputilNewHTTPClient = f
   290  	return func() {
   291  		httputilNewHTTPClient = old
   292  	}
   293  }
   294  
   295  func MockSysconfigConfigureTargetSystem(f func(mod *asserts.Model, opts *sysconfig.Options) error) (restore func()) {
   296  	old := sysconfigConfigureTargetSystem
   297  	sysconfigConfigureTargetSystem = f
   298  	return func() {
   299  		sysconfigConfigureTargetSystem = old
   300  	}
   301  }
   302  
   303  func MockInstallRun(f func(model gadget.Model, gadgetRoot, kernelRoot, device string, options install.Options, observer gadget.ContentObserver, perfTimings timings.Measurer) (*install.InstalledSystemSideData, error)) (restore func()) {
   304  	old := installRun
   305  	installRun = f
   306  	return func() {
   307  		installRun = old
   308  	}
   309  }
   310  
   311  func MockCloudInitStatus(f func() (sysconfig.CloudInitState, error)) (restore func()) {
   312  	old := cloudInitStatus
   313  	cloudInitStatus = f
   314  	return func() {
   315  		cloudInitStatus = old
   316  	}
   317  }
   318  
   319  func MockRestrictCloudInit(f func(sysconfig.CloudInitState, *sysconfig.CloudInitRestrictOptions) (sysconfig.CloudInitRestrictionResult, error)) (restore func()) {
   320  	old := restrictCloudInit
   321  	restrictCloudInit = f
   322  	return func() {
   323  		restrictCloudInit = old
   324  	}
   325  }
   326  
   327  func DeviceManagerHasFDESetupHook(mgr *DeviceManager) (bool, error) {
   328  	return mgr.hasFDESetupHook()
   329  }
   330  
   331  func DeviceManagerRunFDESetupHook(mgr *DeviceManager, req *fde.SetupRequest) ([]byte, error) {
   332  	return mgr.runFDESetupHook(req)
   333  }
   334  
   335  func DeviceManagerCheckEncryption(mgr *DeviceManager, st *state.State, deviceCtx snapstate.DeviceContext) (bool, error) {
   336  	return mgr.checkEncryption(st, deviceCtx)
   337  }
   338  
   339  func DeviceManagerCheckFDEFeatures(mgr *DeviceManager, st *state.State) error {
   340  	return mgr.checkFDEFeatures(st)
   341  }