github.com/hugh712/snapd@v0.0.0-20200910133618-1a99902bd583/overlord/snapstate/snapstatetest/devicectx.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2018-2019 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 snapstatetest
    21  
    22  import (
    23  	"github.com/snapcore/snapd/asserts"
    24  	"github.com/snapcore/snapd/asserts/sysdb"
    25  	"github.com/snapcore/snapd/overlord/snapstate"
    26  	"github.com/snapcore/snapd/overlord/state"
    27  )
    28  
    29  type TrivialDeviceContext struct {
    30  	DeviceModel    *asserts.Model
    31  	OldDeviceModel *asserts.Model
    32  	Remodeling     bool
    33  	CtxStore       snapstate.StoreService
    34  	SysMode        string
    35  	Ground         bool
    36  }
    37  
    38  func (dc *TrivialDeviceContext) Model() *asserts.Model {
    39  	return dc.DeviceModel
    40  }
    41  
    42  func (dc *TrivialDeviceContext) GroundContext() snapstate.DeviceContext {
    43  	if dc.ForRemodeling() && dc.OldDeviceModel != nil {
    44  		return &TrivialDeviceContext{
    45  			DeviceModel: dc.OldDeviceModel,
    46  			SysMode:     dc.SysMode,
    47  			Ground:      true,
    48  		}
    49  	}
    50  	return &TrivialDeviceContext{
    51  		DeviceModel: dc.DeviceModel,
    52  		SysMode:     dc.SysMode,
    53  		Ground:      true,
    54  	}
    55  }
    56  
    57  func (dc *TrivialDeviceContext) Classic() bool {
    58  	return dc.DeviceModel.Classic()
    59  }
    60  
    61  func (dc *TrivialDeviceContext) Kernel() string {
    62  	return dc.DeviceModel.Kernel()
    63  }
    64  
    65  func (dc *TrivialDeviceContext) Base() string {
    66  	return dc.DeviceModel.Base()
    67  }
    68  
    69  func (dc *TrivialDeviceContext) HasModeenv() bool {
    70  	return dc.Model().Grade() != asserts.ModelGradeUnset
    71  }
    72  
    73  func (dc *TrivialDeviceContext) RunMode() bool {
    74  	return dc.SystemMode() == "run"
    75  }
    76  
    77  func (dc *TrivialDeviceContext) Store() snapstate.StoreService {
    78  	if dc.Ground {
    79  		panic("retrieved ground context is not intended to drive store operations")
    80  	}
    81  	return dc.CtxStore
    82  }
    83  
    84  func (dc *TrivialDeviceContext) ForRemodeling() bool {
    85  	return dc.Remodeling
    86  }
    87  
    88  func (dc *TrivialDeviceContext) SystemMode() string {
    89  	mode := dc.SysMode
    90  	if mode == "" {
    91  		return "run"
    92  	}
    93  	return mode
    94  }
    95  
    96  func MockDeviceModel(model *asserts.Model) (restore func()) {
    97  	var deviceCtx snapstate.DeviceContext
    98  	if model != nil {
    99  		deviceCtx = &TrivialDeviceContext{DeviceModel: model}
   100  	}
   101  	return MockDeviceContext(deviceCtx)
   102  }
   103  
   104  func MockDeviceModelAndMode(model *asserts.Model, systemMode string) (restore func()) {
   105  	deviceCtx := &TrivialDeviceContext{DeviceModel: model, SysMode: systemMode}
   106  	return MockDeviceContext(deviceCtx)
   107  }
   108  
   109  func MockDeviceContext(deviceCtx snapstate.DeviceContext) (restore func()) {
   110  	deviceCtxHook := func(st *state.State, task *state.Task, providedDeviceCtx snapstate.DeviceContext) (snapstate.DeviceContext, error) {
   111  		if providedDeviceCtx != nil {
   112  			return providedDeviceCtx, nil
   113  		}
   114  		if deviceCtx == nil {
   115  			return nil, state.ErrNoState
   116  		}
   117  		return deviceCtx, nil
   118  	}
   119  	r1 := ReplaceDeviceCtxHook(deviceCtxHook)
   120  	// for convenience reflect from the context whether there is a
   121  	// remodeling
   122  	r2 := ReplaceRemodelingHook(func(*state.State) bool {
   123  		return deviceCtx != nil && deviceCtx.ForRemodeling()
   124  	})
   125  	return func() {
   126  		r1()
   127  		r2()
   128  	}
   129  }
   130  
   131  func ReplaceDeviceCtxHook(deviceCtxHook func(st *state.State, task *state.Task, providedDeviceCtx snapstate.DeviceContext) (snapstate.DeviceContext, error)) (restore func()) {
   132  	oldHook := snapstate.DeviceCtx
   133  	snapstate.DeviceCtx = deviceCtxHook
   134  	return func() {
   135  		snapstate.DeviceCtx = oldHook
   136  	}
   137  }
   138  
   139  func UseFallbackDeviceModel() (restore func()) {
   140  	return MockDeviceModel(sysdb.GenericClassicModel())
   141  }
   142  
   143  func ReplaceRemodelingHook(remodelingHook func(st *state.State) bool) (restore func()) {
   144  	oldHook := snapstate.Remodeling
   145  	snapstate.Remodeling = remodelingHook
   146  	return func() {
   147  		snapstate.Remodeling = oldHook
   148  	}
   149  }