github.com/hugh712/snapd@v0.0.0-20200910133618-1a99902bd583/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) asserts.KeypairManager {
    84  	return m.keypairMgr
    85  }
    86  
    87  func EnsureOperationalShouldBackoff(m *DeviceManager, now time.Time) bool {
    88  	return m.ensureOperationalShouldBackoff(now)
    89  }
    90  
    91  func BecomeOperationalBackoff(m *DeviceManager) time.Duration {
    92  	return m.becomeOperationalBackoff
    93  }
    94  
    95  func SetLastBecomeOperationalAttempt(m *DeviceManager, t time.Time) {
    96  	m.lastBecomeOperationalAttempt = t
    97  }
    98  
    99  func SetSystemMode(m *DeviceManager, mode string) {
   100  	m.systemMode = mode
   101  }
   102  
   103  func SetTimeOnce(m *DeviceManager, name string, t time.Time) error {
   104  	return m.setTimeOnce(name, t)
   105  }
   106  
   107  func MockRepeatRequestSerial(label string) (restore func()) {
   108  	old := repeatRequestSerial
   109  	repeatRequestSerial = label
   110  	return func() {
   111  		repeatRequestSerial = old
   112  	}
   113  }
   114  
   115  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()) {
   116  	old := snapstateInstallWithDeviceContext
   117  	snapstateInstallWithDeviceContext = f
   118  	return func() {
   119  		snapstateInstallWithDeviceContext = old
   120  	}
   121  }
   122  
   123  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()) {
   124  	old := snapstateUpdateWithDeviceContext
   125  	snapstateUpdateWithDeviceContext = f
   126  	return func() {
   127  		snapstateUpdateWithDeviceContext = old
   128  	}
   129  }
   130  
   131  func EnsureSeeded(m *DeviceManager) error {
   132  	return m.ensureSeeded()
   133  }
   134  
   135  func EnsureCloudInitRestricted(m *DeviceManager) error {
   136  	return m.ensureCloudInitRestricted()
   137  }
   138  
   139  var PopulateStateFromSeedImpl = populateStateFromSeedImpl
   140  
   141  type PopulateStateFromSeedOptions = populateStateFromSeedOptions
   142  
   143  func MockPopulateStateFromSeed(f func(*state.State, *PopulateStateFromSeedOptions, timings.Measurer) ([]*state.TaskSet, error)) (restore func()) {
   144  	old := populateStateFromSeed
   145  	populateStateFromSeed = f
   146  	return func() {
   147  		populateStateFromSeed = old
   148  	}
   149  }
   150  
   151  func EnsureBootOk(m *DeviceManager) error {
   152  	return m.ensureBootOk()
   153  }
   154  
   155  func SetBootOkRan(m *DeviceManager, b bool) {
   156  	m.bootOkRan = b
   157  }
   158  
   159  func StartTime() time.Time {
   160  	return startTime
   161  }
   162  
   163  type (
   164  	RegistrationContext = registrationContext
   165  	RemodelContext      = remodelContext
   166  	SeededSystem        = seededSystem
   167  )
   168  
   169  func RegistrationCtx(m *DeviceManager, t *state.Task) (registrationContext, error) {
   170  	return m.registrationCtx(t)
   171  }
   172  
   173  func RemodelDeviceBackend(remodCtx remodelContext) storecontext.DeviceBackend {
   174  	return remodCtx.(interface {
   175  		deviceBackend() storecontext.DeviceBackend
   176  	}).deviceBackend()
   177  }
   178  
   179  var (
   180  	ImportAssertionsFromSeed     = importAssertionsFromSeed
   181  	CheckGadgetOrKernel          = checkGadgetOrKernel
   182  	CheckGadgetValid             = checkGadgetValid
   183  	CheckGadgetRemodelCompatible = checkGadgetRemodelCompatible
   184  	CanAutoRefresh               = canAutoRefresh
   185  	NewEnoughProxy               = newEnoughProxy
   186  
   187  	IncEnsureOperationalAttempts = incEnsureOperationalAttempts
   188  	EnsureOperationalAttempts    = ensureOperationalAttempts
   189  
   190  	RemodelTasks = remodelTasks
   191  
   192  	RemodelCtx        = remodelCtx
   193  	CleanupRemodelCtx = cleanupRemodelCtx
   194  	CachedRemodelCtx  = cachedRemodelCtx
   195  
   196  	GadgetUpdateBlocked = gadgetUpdateBlocked
   197  	CurrentGadgetInfo   = currentGadgetInfo
   198  	PendingGadgetInfo   = pendingGadgetInfo
   199  
   200  	CriticalTaskEdges = criticalTaskEdges
   201  )
   202  
   203  func MockGadgetUpdate(mock func(current, update gadget.GadgetData, path string, policy gadget.UpdatePolicyFunc, observer gadget.ContentUpdateObserver) error) (restore func()) {
   204  	old := gadgetUpdate
   205  	gadgetUpdate = mock
   206  	return func() {
   207  		gadgetUpdate = old
   208  	}
   209  }
   210  
   211  func MockGadgetIsCompatible(mock func(current, update *gadget.Info) error) (restore func()) {
   212  	old := gadgetIsCompatible
   213  	gadgetIsCompatible = mock
   214  	return func() {
   215  		gadgetIsCompatible = old
   216  	}
   217  }
   218  
   219  func MockBootMakeBootable(f func(model *asserts.Model, rootdir string, bootWith *boot.BootableSet, seal *boot.TrustedAssetsInstallObserver) error) (restore func()) {
   220  	old := bootMakeBootable
   221  	bootMakeBootable = f
   222  	return func() {
   223  		bootMakeBootable = old
   224  	}
   225  }
   226  
   227  func MockSecbootCheckKeySealingSupported(f func() error) (restore func()) {
   228  	old := secbootCheckKeySealingSupported
   229  	secbootCheckKeySealingSupported = f
   230  	return func() {
   231  		secbootCheckKeySealingSupported = old
   232  	}
   233  }
   234  
   235  func MockHttputilNewHTTPClient(f func(opts *httputil.ClientOptions) *http.Client) (restore func()) {
   236  	old := httputilNewHTTPClient
   237  	httputilNewHTTPClient = f
   238  	return func() {
   239  		httputilNewHTTPClient = old
   240  	}
   241  }
   242  
   243  func MockSysconfigConfigureRunSystem(f func(opts *sysconfig.Options) error) (restore func()) {
   244  	old := sysconfigConfigureRunSystem
   245  	sysconfigConfigureRunSystem = f
   246  	return func() {
   247  		sysconfigConfigureRunSystem = old
   248  	}
   249  }
   250  
   251  func MockInstallRun(f func(gadgetRoot, device string, options install.Options, observer install.SystemInstallObserver) error) (restore func()) {
   252  	old := installRun
   253  	installRun = f
   254  	return func() {
   255  		installRun = old
   256  	}
   257  }
   258  
   259  func MockCloudInitStatus(f func() (sysconfig.CloudInitState, error)) (restore func()) {
   260  	old := cloudInitStatus
   261  	cloudInitStatus = f
   262  	return func() {
   263  		cloudInitStatus = old
   264  	}
   265  }
   266  
   267  func MockRestrictCloudInit(f func(sysconfig.CloudInitState, *sysconfig.CloudInitRestrictOptions) (sysconfig.CloudInitRestrictionResult, error)) (restore func()) {
   268  	old := restrictCloudInit
   269  	restrictCloudInit = f
   270  	return func() {
   271  		restrictCloudInit = old
   272  	}
   273  }