github.com/stolowski/snapd@v0.0.0-20210407085831-115137ce5a22/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/overlord/snapstate"
    33  	"github.com/snapcore/snapd/overlord/state"
    34  	"github.com/snapcore/snapd/overlord/storecontext"
    35  	"github.com/snapcore/snapd/sysconfig"
    36  	"github.com/snapcore/snapd/timings"
    37  )
    38  
    39  func MockKeyLength(n int) (restore func()) {
    40  	if n < 1024 {
    41  		panic("key length must be >= 1024")
    42  	}
    43  
    44  	oldKeyLength := keyLength
    45  	keyLength = n
    46  	return func() {
    47  		keyLength = oldKeyLength
    48  	}
    49  }
    50  
    51  func MockBaseStoreURL(url string) (restore func()) {
    52  	oldURL := baseStoreURL
    53  	baseStoreURL = mustParse(url).ResolveReference(authRef)
    54  	return func() {
    55  		baseStoreURL = oldURL
    56  	}
    57  }
    58  
    59  func MockRetryInterval(interval time.Duration) (restore func()) {
    60  	old := retryInterval
    61  	retryInterval = interval
    62  	return func() {
    63  		retryInterval = old
    64  	}
    65  }
    66  
    67  func MockMaxTentatives(max int) (restore func()) {
    68  	old := maxTentatives
    69  	maxTentatives = max
    70  	return func() {
    71  		maxTentatives = old
    72  	}
    73  }
    74  
    75  func MockTimeNow(f func() time.Time) (restore func()) {
    76  	old := timeNow
    77  	timeNow = f
    78  	return func() {
    79  		timeNow = old
    80  	}
    81  }
    82  
    83  func KeypairManager(m *DeviceManager) (keypairMgr asserts.KeypairManager) {
    84  	// XXX expose the with... method at some point
    85  	err := m.withKeypairMgr(func(km asserts.KeypairManager) error {
    86  		keypairMgr = km
    87  		return nil
    88  	})
    89  	if err != nil {
    90  		panic(err)
    91  	}
    92  	return keypairMgr
    93  }
    94  
    95  func SaveAvailable(m *DeviceManager) bool {
    96  	return m.saveAvailable
    97  }
    98  
    99  func SetSaveAvailable(m *DeviceManager, avail bool) {
   100  	m.saveAvailable = avail
   101  }
   102  
   103  func EnsureOperationalShouldBackoff(m *DeviceManager, now time.Time) bool {
   104  	return m.ensureOperationalShouldBackoff(now)
   105  }
   106  
   107  func BecomeOperationalBackoff(m *DeviceManager) time.Duration {
   108  	return m.becomeOperationalBackoff
   109  }
   110  
   111  func SetLastBecomeOperationalAttempt(m *DeviceManager, t time.Time) {
   112  	m.lastBecomeOperationalAttempt = t
   113  }
   114  
   115  func SetSystemMode(m *DeviceManager, mode string) {
   116  	m.systemMode = mode
   117  }
   118  
   119  func SetTimeOnce(m *DeviceManager, name string, t time.Time) error {
   120  	return m.setTimeOnce(name, t)
   121  }
   122  
   123  func PreloadGadget(m *DeviceManager) (*gadget.Info, error) {
   124  	return m.preloadGadget()
   125  }
   126  
   127  func MockRepeatRequestSerial(label string) (restore func()) {
   128  	old := repeatRequestSerial
   129  	repeatRequestSerial = label
   130  	return func() {
   131  		repeatRequestSerial = old
   132  	}
   133  }
   134  
   135  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()) {
   136  	old := snapstateInstallWithDeviceContext
   137  	snapstateInstallWithDeviceContext = f
   138  	return func() {
   139  		snapstateInstallWithDeviceContext = old
   140  	}
   141  }
   142  
   143  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()) {
   144  	old := snapstateUpdateWithDeviceContext
   145  	snapstateUpdateWithDeviceContext = f
   146  	return func() {
   147  		snapstateUpdateWithDeviceContext = old
   148  	}
   149  }
   150  
   151  func EnsureSeeded(m *DeviceManager) error {
   152  	return m.ensureSeeded()
   153  }
   154  
   155  func EnsureCloudInitRestricted(m *DeviceManager) error {
   156  	return m.ensureCloudInitRestricted()
   157  }
   158  
   159  var PopulateStateFromSeedImpl = populateStateFromSeedImpl
   160  
   161  type PopulateStateFromSeedOptions = populateStateFromSeedOptions
   162  
   163  func MockPopulateStateFromSeed(f func(*state.State, *PopulateStateFromSeedOptions, timings.Measurer) ([]*state.TaskSet, error)) (restore func()) {
   164  	old := populateStateFromSeed
   165  	populateStateFromSeed = f
   166  	return func() {
   167  		populateStateFromSeed = old
   168  	}
   169  }
   170  
   171  func EnsureBootOk(m *DeviceManager) error {
   172  	return m.ensureBootOk()
   173  }
   174  
   175  func SetBootOkRan(m *DeviceManager, b bool) {
   176  	m.bootOkRan = b
   177  }
   178  
   179  func SetInstalledRan(m *DeviceManager, b bool) {
   180  	m.ensureInstalledRan = b
   181  }
   182  
   183  func StartTime() time.Time {
   184  	return startTime
   185  }
   186  
   187  type (
   188  	RegistrationContext = registrationContext
   189  	RemodelContext      = remodelContext
   190  	SeededSystem        = seededSystem
   191  )
   192  
   193  func RegistrationCtx(m *DeviceManager, t *state.Task) (registrationContext, error) {
   194  	return m.registrationCtx(t)
   195  }
   196  
   197  func RemodelDeviceBackend(remodCtx remodelContext) storecontext.DeviceBackend {
   198  	return remodCtx.(interface {
   199  		deviceBackend() storecontext.DeviceBackend
   200  	}).deviceBackend()
   201  }
   202  
   203  var (
   204  	LoadDeviceSeed               = loadDeviceSeed
   205  	UnloadDeviceSeed             = unloadDeviceSeed
   206  	ImportAssertionsFromSeed     = importAssertionsFromSeed
   207  	CheckGadgetOrKernel          = checkGadgetOrKernel
   208  	CheckGadgetValid             = checkGadgetValid
   209  	CheckGadgetRemodelCompatible = checkGadgetRemodelCompatible
   210  	CanAutoRefresh               = canAutoRefresh
   211  	NewEnoughProxy               = newEnoughProxy
   212  
   213  	IncEnsureOperationalAttempts = incEnsureOperationalAttempts
   214  	EnsureOperationalAttempts    = ensureOperationalAttempts
   215  
   216  	RemodelTasks = remodelTasks
   217  
   218  	RemodelCtx        = remodelCtx
   219  	CleanupRemodelCtx = cleanupRemodelCtx
   220  	CachedRemodelCtx  = cachedRemodelCtx
   221  
   222  	GadgetUpdateBlocked = gadgetUpdateBlocked
   223  	CurrentGadgetInfo   = currentGadgetInfo
   224  	PendingGadgetInfo   = pendingGadgetInfo
   225  
   226  	CriticalTaskEdges = criticalTaskEdges
   227  )
   228  
   229  func MockGadgetUpdate(mock func(current, update gadget.GadgetData, path string, policy gadget.UpdatePolicyFunc, observer gadget.ContentUpdateObserver) error) (restore func()) {
   230  	old := gadgetUpdate
   231  	gadgetUpdate = mock
   232  	return func() {
   233  		gadgetUpdate = old
   234  	}
   235  }
   236  
   237  func MockGadgetIsCompatible(mock func(current, update *gadget.Info) error) (restore func()) {
   238  	old := gadgetIsCompatible
   239  	gadgetIsCompatible = mock
   240  	return func() {
   241  		gadgetIsCompatible = old
   242  	}
   243  }
   244  
   245  func MockBootMakeSystemRunnable(f func(model *asserts.Model, bootWith *boot.BootableSet, seal *boot.TrustedAssetsInstallObserver) error) (restore func()) {
   246  	old := bootMakeRunnable
   247  	bootMakeRunnable = f
   248  	return func() {
   249  		bootMakeRunnable = old
   250  	}
   251  }
   252  
   253  func MockBootEnsureNextBootToRunMode(f func(systemLabel string) error) (restore func()) {
   254  	old := bootEnsureNextBootToRunMode
   255  	bootEnsureNextBootToRunMode = f
   256  	return func() {
   257  		bootEnsureNextBootToRunMode = old
   258  	}
   259  }
   260  
   261  func MockSecbootCheckKeySealingSupported(f func() error) (restore func()) {
   262  	old := secbootCheckKeySealingSupported
   263  	secbootCheckKeySealingSupported = f
   264  	return func() {
   265  		secbootCheckKeySealingSupported = old
   266  	}
   267  }
   268  
   269  func MockHttputilNewHTTPClient(f func(opts *httputil.ClientOptions) *http.Client) (restore func()) {
   270  	old := httputilNewHTTPClient
   271  	httputilNewHTTPClient = f
   272  	return func() {
   273  		httputilNewHTTPClient = old
   274  	}
   275  }
   276  
   277  func MockSysconfigConfigureTargetSystem(f func(opts *sysconfig.Options) error) (restore func()) {
   278  	old := sysconfigConfigureTargetSystem
   279  	sysconfigConfigureTargetSystem = f
   280  	return func() {
   281  		sysconfigConfigureTargetSystem = old
   282  	}
   283  }
   284  
   285  func MockInstallRun(f func(model gadget.Model, gadgetRoot, kernelRoot, device string, options install.Options, observer gadget.ContentObserver) (*install.InstalledSystemSideData, error)) (restore func()) {
   286  	old := installRun
   287  	installRun = f
   288  	return func() {
   289  		installRun = old
   290  	}
   291  }
   292  
   293  func MockCloudInitStatus(f func() (sysconfig.CloudInitState, error)) (restore func()) {
   294  	old := cloudInitStatus
   295  	cloudInitStatus = f
   296  	return func() {
   297  		cloudInitStatus = old
   298  	}
   299  }
   300  
   301  func MockRestrictCloudInit(f func(sysconfig.CloudInitState, *sysconfig.CloudInitRestrictOptions) (sysconfig.CloudInitRestrictionResult, error)) (restore func()) {
   302  	old := restrictCloudInit
   303  	restrictCloudInit = f
   304  	return func() {
   305  		restrictCloudInit = old
   306  	}
   307  }
   308  
   309  func DeviceManagerHasFDESetupHook(mgr *DeviceManager) (bool, error) {
   310  	return mgr.hasFDESetupHook()
   311  }
   312  
   313  func DeviceManagerRunFDESetupHook(mgr *DeviceManager, op string, params *boot.FDESetupHookParams) ([]byte, error) {
   314  	return mgr.runFDESetupHook(op, params)
   315  }
   316  
   317  func DeviceManagerCheckEncryption(mgr *DeviceManager, st *state.State, deviceCtx snapstate.DeviceContext) (bool, error) {
   318  	return mgr.checkEncryption(st, deviceCtx)
   319  }
   320  
   321  func DeviceManagerCheckFDEFeatures(mgr *DeviceManager, st *state.State) error {
   322  	return mgr.checkFDEFeatures(st)
   323  }