// -*- Mode: Go; indent-tabs-mode: t -*-

/*
 * Copyright (C) 2019 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package snapstate_test

import (
	. "gopkg.in/check.v1"

	"github.com/snapcore/snapd/dirs"
	"github.com/snapcore/snapd/overlord/configstate/config"
	"github.com/snapcore/snapd/overlord/dot/dottest"
	"github.com/snapcore/snapd/overlord/restart"
	"github.com/snapcore/snapd/overlord/snapstate"
	"github.com/snapcore/snapd/overlord/snapstate/snapstatetest"
	"github.com/snapcore/snapd/overlord/state"
	"github.com/snapcore/snapd/snap"
	"github.com/snapcore/snapd/testutil"
)

type rebootSuite struct {
	testutil.BaseTest
	state *state.State
}

var _ = Suite(&rebootSuite{})

func (s *rebootSuite) SetUpTest(c *C) {
	s.BaseTest.SetUpTest(c)
	dirs.SetRootDir(c.MkDir())
	s.AddCleanup(func() { dirs.SetRootDir("") })
	s.state = state.New(nil)
}

func (s *rebootSuite) TearDownTest(c *C) {
	restore := dottest.RegisterChangeExporter(c, s.state)
	s.BaseTest.AddCleanup(restore)
	s.BaseTest.TearDownTest(c)
}

func (s *rebootSuite) deviceCtx(c *C) snapstate.DeviceContext {
	dctx, err := snapstate.DeviceCtx(s.state, nil, nil)
	c.Assert(err, IsNil)
	return dctx
}

func (s *rebootSuite) snapInstallTaskSetForSnapSetup(snapName, base string, snapType snap.Type) snapstate.SnapInstallTaskSet {
	snapsup := &snapstate.SnapSetup{
		SideInfo: &snap.SideInfo{
			RealName: snapName,
			SnapID:   snapName,
			Revision: snap.R(1),
		},
		Type: snapType,
		Base: base,
	}
	prereq := s.state.NewTask("prerequisites", "...")
	prereq.Set("snap-setup", snapsup)
	prepareSnap := s.state.NewTask("prepare-snap", "...")
	prepareSnap.WaitFor(prereq)
	mountSnap := s.state.NewTask("mount-snap", "...")
	mountSnap.WaitFor(prepareSnap)
	unlinkSnap := s.state.NewTask("unlink-snap", "...")
	unlinkSnap.WaitFor(mountSnap)
	linkSnap := s.state.NewTask("link-snap", "...")
	linkSnap.WaitFor(unlinkSnap)
	autoConnect := s.state.NewTask("auto-connect", "...")
	autoConnect.WaitFor(linkSnap)
	startServices := s.state.NewTask("start-snap-services", "...")
	startServices.WaitFor(autoConnect)
	ts := state.NewTaskSet(prereq, prepareSnap, mountSnap, unlinkSnap, linkSnap, autoConnect, startServices)

	ts.MarkEdge(prereq, snapstate.BeginEdge)
	ts.MarkEdge(prepareSnap, snapstate.LastBeforeLocalModificationsEdge)
	ts.MarkEdge(linkSnap, snapstate.MaybeRebootEdge)
	ts.MarkEdge(autoConnect, snapstate.MaybeRebootWaitEdge)
	ts.MarkEdge(startServices, snapstate.EndEdge)

	ts.JoinLane(s.state.NewLane())

	return snapstate.NewSnapInstallTaskSetForTest(
		snapsup,
		ts,
		[]*state.Task{prereq, prepareSnap}, // before local modification tasks
		mountSnap,
		[]*state.Task{unlinkSnap, linkSnap}, // modification inducing tasks before reboot
		[]*state.Task{autoConnect, startServices}, // post reboot tasks
	)
}

func taskSetsFromInstallSets(stss []snapstate.SnapInstallTaskSet) []*state.TaskSet {
	tss := make([]*state.TaskSet, 0, len(stss))
	for _, sts := range stss {
		tss = append(tss, sts.TaskSet())
	}
	return tss
}

func taskSetLanes(ts *state.TaskSet) []int {
	var lanes []int
	seen := make(map[int]bool)
	for _, t := range ts.Tasks() {
		for _, l := range t.Lanes() {
			if seen[l] {
				continue
			}
			seen[l] = true
			lanes = append(lanes, l)
		}
	}
	return lanes
}

func (s *rebootSuite) TestTaskSetsByTypeForEssentialSnapsNoBootBase(c *C) {
	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("my-base", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("my-os", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	mappedTaskSets, err := snapstate.TaskSetsByTypeForEssentialSnaps(taskSetsFromInstallSets(stss), "")
	c.Assert(err, IsNil)
	c.Check(mappedTaskSets, DeepEquals, map[snap.Type]*state.TaskSet{
		snap.TypeGadget: stss[1].TaskSet(),
		snap.TypeKernel: stss[2].TaskSet(),
	})
}

func (s *rebootSuite) TestTaskSetsByTypeForEssentialSnapsBootBase(c *C) {
	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("my-base", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("my-os", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	mappedTaskSets, err := snapstate.TaskSetsByTypeForEssentialSnaps(taskSetsFromInstallSets(stss), "my-base")
	c.Assert(err, IsNil)
	c.Check(mappedTaskSets, DeepEquals, map[snap.Type]*state.TaskSet{
		snap.TypeBase:   stss[0].TaskSet(),
		snap.TypeGadget: stss[1].TaskSet(),
		snap.TypeKernel: stss[2].TaskSet(),
	})
}

func (s *rebootSuite) TestSetDefaultRestartBoundariesLink(c *C) {
	s.state.Lock()
	defer s.state.Unlock()

	t1 := s.state.NewTask("first-thing", "...")
	t2 := s.state.NewTask("other-thing", "...")
	ts := state.NewTaskSet(t1, t2)
	ts.MarkEdge(t1, snapstate.MaybeRebootEdge)

	snapstate.SetDefaultRestartBoundaries(ts)

	var boundary restart.RestartBoundaryDirection
	c.Check(t1.Get("restart-boundary", &boundary), IsNil)
	c.Check(t2.Get("restart-boundary", &boundary), ErrorMatches, `no state entry for key "restart-boundary"`)
}

func (s *rebootSuite) TestSetDefaultRestartBoundariesUnlink(c *C) {
	s.state.Lock()
	defer s.state.Unlock()

	t1 := s.state.NewTask("unlink-snap", "...")
	t2 := s.state.NewTask("other-thing", "...")
	ts := state.NewTaskSet(t1, t2)

	snapstate.SetDefaultRestartBoundaries(ts)

	var boundary restart.RestartBoundaryDirection
	c.Check(t1.Get("restart-boundary", &boundary), IsNil)
	c.Check(t2.Get("restart-boundary", &boundary), ErrorMatches, `no state entry for key "restart-boundary"`)
}

func (s *rebootSuite) TestSetDefaultRestartBoundariesUnlinkCurrent(c *C) {
	s.state.Lock()
	defer s.state.Unlock()

	t1 := s.state.NewTask("unlink-current-snap", "...")
	t2 := s.state.NewTask("other-thing", "...")
	ts := state.NewTaskSet(t1, t2)

	snapstate.SetDefaultRestartBoundaries(ts)

	var boundary restart.RestartBoundaryDirection
	c.Check(t1.Get("restart-boundary", &boundary), IsNil)
	c.Check(t2.Get("restart-boundary", &boundary), ErrorMatches, `no state entry for key "restart-boundary"`)
}

func (s *rebootSuite) TestDeviceModelBootBaseEmpty(c *C) {
	defer snapstatetest.MockDeviceModel(nil)()
	bootBase, err := snapstate.DeviceModelBootBase(s.state, nil)
	c.Check(err, IsNil)
	c.Check(bootBase, Equals, "")
}

func (s *rebootSuite) TestDeviceModelBootBaseUC16(c *C) {
	defer snapstatetest.MockDeviceModel(DefaultModel())()
	bootBase, err := snapstate.DeviceModelBootBase(s.state, nil)
	c.Check(err, IsNil)
	c.Check(bootBase, Equals, "core")
}

func (s *rebootSuite) TestDeviceModelBootBaseUC(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()
	bootBase, err := snapstate.DeviceModelBootBase(s.state, nil)
	c.Check(err, IsNil)
	c.Check(bootBase, Equals, "core20")
}

func (s *rebootSuite) TestDeviceModelBootBaseClassic(c *C) {
	defer snapstatetest.MockDeviceModel(ClassicModel())()
	bootBase, err := snapstate.DeviceModelBootBase(s.state, nil)
	c.Check(err, IsNil)
	c.Check(bootBase, Equals, "core")
}

func (s *rebootSuite) TestDeviceModelBootBaseClassicModelProvided(c *C) {
	// Mock the model to classic
	defer snapstatetest.MockDeviceModel(ClassicModel())()

	// then provide a UC18 model instead
	deviceCtx := &snapstatetest.TrivialDeviceContext{DeviceModel: ModelWithBase("core18")}
	bootBase, err := snapstate.DeviceModelBootBase(s.state, deviceCtx)
	c.Check(err, IsNil)
	c.Check(bootBase, Equals, "core18")
}

func findUnlinkTask(ts *state.TaskSet) *state.Task {
	for _, t := range ts.Tasks() {
		switch t.Kind() {
		case "unlink-snap", "unlink-current-snap":
			return t
		}
	}
	return nil
}

func findTaskKind(ts *state.TaskSet, kind string) *state.Task {
	for _, t := range ts.Tasks() {
		if t.Kind() == kind {
			return t
		}
	}
	return nil
}

func (s *rebootSuite) hasRestartBoundaries(c *C, ts *state.TaskSet) bool {
	t1 := ts.MaybeEdge(snapstate.MaybeRebootEdge)
	t2 := findUnlinkTask(ts)
	c.Assert(t1, NotNil)
	c.Assert(t2, NotNil)

	var boundary restart.RestartBoundaryDirection
	if err := t1.Get("restart-boundary", &boundary); err != nil {
		return false
	}
	if err := t2.Get("restart-boundary", &boundary); err != nil {
		return false
	}
	return true
}

func (s *rebootSuite) hasDoRestartBoundaries(c *C, ts *state.TaskSet) bool {
	t := ts.MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(t, NotNil)

	var boundary restart.RestartBoundaryDirection
	if err := t.Get("restart-boundary", &boundary); err != nil {
		return false
	}
	return true
}

func (s *rebootSuite) hasUndoRestartBoundaries(c *C, ts *state.TaskSet) bool {
	t := findUnlinkTask(ts)
	c.Assert(t, NotNil)

	var boundary restart.RestartBoundaryDirection
	if err := t.Get("restart-boundary", &boundary); err != nil {
		return false
	}
	return true
}

func (s *rebootSuite) TestSetEssentialSnapsRestartBoundariesUC16(c *C) {
	defer snapstatetest.MockDeviceModel(DefaultModel())()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("core", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	err := snapstate.SetEssentialSnapsRestartBoundaries(s.state, nil, taskSetsFromInstallSets(stss))
	c.Assert(err, IsNil)
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[3].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[4].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestSetEssentialSnapsRestartBoundariesUC20(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("core", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	err := snapstate.SetEssentialSnapsRestartBoundaries(s.state, nil, taskSetsFromInstallSets(stss))
	c.Assert(err, IsNil)
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[3].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[4].TaskSet()), Equals, false)
}

func (s *rebootSuite) setDependsOn(c *C, ts, dep *state.TaskSet) bool {
	firstTaskOfTS, err := ts.Edge(snapstate.BeginEdge)
	c.Assert(err, IsNil)
	lastTaskOfDep, err := dep.Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	for _, wt := range firstTaskOfTS.WaitTasks() {
		if wt == lastTaskOfDep {
			return true
		}
	}
	return false
}

func waitsOnTransitively(waiter, target *state.Task) bool {
	stack := append([]*state.Task(nil), waiter.WaitTasks()...)
	seen := make(map[string]bool, len(stack))
	for len(stack) > 0 {
		cur := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if cur == target {
			return true
		}
		if seen[cur.ID()] {
			continue
		}
		seen[cur.ID()] = true
		stack = append(stack, cur.WaitTasks()...)
	}

	return false
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsUC16NoSplits(c *C) {
	defer snapstatetest.MockDeviceModel(DefaultModel())()

	s.state.Lock()
	defer s.state.Unlock()

	// Run without gadget, as that will make it also non-split currently
	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// core, kernel should have individual restart boundaries
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[3].TaskSet()), Equals, false)

	// core and kernel are not transactional on UC16
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsSnapdAndEssential(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snapd", "", snap.TypeSnapd),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Snapd should have no restart boundaries
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, false)

	// boot-base, gadget, kernel setup for single-reboot
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[3].TaskSet()), Equals, true)
	c.Check(s.hasUndoRestartBoundaries(c, stss[3].TaskSet()), Equals, false)

	// TypeApp should have no boundaries
	c.Check(s.hasRestartBoundaries(c, stss[4].TaskSet()), Equals, false)

	// snapd is refreshed in this change. thus, essential snaps should not start
	// prerequisites/download before snapd fully finishes.
	snapdEndTask, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)
	for _, idx := range []int{1, 2, 3} {
		beginTask, err := stss[idx].TaskSet().Edge(snapstate.BeginEdge)
		c.Assert(err, IsNil)
		c.Check(beginTask.WaitTasks(), testutil.Contains, snapdEndTask)
	}

	// base, gadget and kernel are transactional
	c.Check(taskSetsShareLane(stss[1].TaskSet(), stss[2].TaskSet(), stss[3].TaskSet()), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsBaseKernel(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Expect restart boundaries on both
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)

	linkSnapBase := stss[0].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapBase, NotNil)
	mountSnapBase := findTaskKind(stss[0].TaskSet(), "mount-snap")
	c.Assert(mountSnapBase, NotNil)
	unlinkSnapBase := findUnlinkTask(stss[0].TaskSet())
	c.Assert(unlinkSnapBase, NotNil)
	linkSnapKernel := stss[1].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapKernel, NotNil)
	mountSnapKernel := findTaskKind(stss[1].TaskSet(), "mount-snap")
	c.Assert(mountSnapKernel, NotNil)
	unlinkSnapKernel := findUnlinkTask(stss[1].TaskSet())
	c.Assert(unlinkSnapKernel, NotNil)

	// Grab the tasks we need to check dependencies between
	firstTaskOfKernel := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	beginTaskOfKernel, err := stss[1].TaskSet().Edge(snapstate.BeginEdge)
	c.Assert(err, IsNil)
	linkTaskOfKernel, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfKernel, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	linkTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	lastTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	// Things that must be correct:
	// - first local modification task of kernel must be its mount task and must depend on the base mount
	c.Check(firstTaskOfKernel, Equals, mountSnapKernel)
	c.Check(firstTaskOfKernel.WaitTasks(), testutil.Contains, mountSnapBase)
	// - the base's remaining pre-reboot work only starts after the kernel mount phase finishes
	c.Check(unlinkSnapBase.WaitTasks(), testutil.Contains, mountSnapKernel)
	// - the first post-mount task of kernel must depend on the base link
	c.Check(unlinkSnapKernel.WaitTasks(), testutil.Contains, linkTaskOfBase)
	// - prerequisites/download should not be serialized behind base link
	c.Check(beginTaskOfKernel.WaitTasks(), Not(testutil.Contains), linkTaskOfBase)
	// - "auto-connect" (MaybeRebootWaitEdge) of base must depend on "link-snap" of kernel (MaybeRebootEdge)
	c.Check(acTaskOfBase.WaitTasks(), testutil.Contains, linkTaskOfKernel)
	// - "auto-connect" (MaybeRebootWaitEdge) of kernel must depend on the last task of base (EndEdge)
	c.Check(acTaskOfKernel.WaitTasks(), testutil.Contains, lastTaskOfBase)

	// both should be transactional
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, true)

	// Since they are set up for single-reboot, the base should have restart
	// boundaries for the undo path, and kernel should have for do path.
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsBaseGadget(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Expect restart boundaries on both
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)

	linkSnapBase := stss[0].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapBase, NotNil)
	mountSnapBase := findTaskKind(stss[0].TaskSet(), "mount-snap")
	c.Assert(mountSnapBase, NotNil)
	unlinkSnapBase := findUnlinkTask(stss[0].TaskSet())
	c.Assert(unlinkSnapBase, NotNil)
	linkSnapGadget := stss[1].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapGadget, NotNil)
	mountSnapGadget := findTaskKind(stss[1].TaskSet(), "mount-snap")
	c.Assert(mountSnapGadget, NotNil)
	unlinkSnapGadget := findUnlinkTask(stss[1].TaskSet())
	c.Assert(unlinkSnapGadget, NotNil)

	// Grab the tasks we need to check dependencies between
	firstTaskOfGadget := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	linkTaskOfGadget, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfGadget, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	linkTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	lastTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	// Things that must be correct:
	// - first local modification task of gadget must be its mount task and must depend on the base mount
	c.Check(firstTaskOfGadget, Equals, mountSnapGadget)
	c.Check(firstTaskOfGadget.WaitTasks(), testutil.Contains, mountSnapBase)
	// - the base's remaining pre-reboot work only starts after the gadget mount phase finishes
	c.Check(unlinkSnapBase.WaitTasks(), testutil.Contains, mountSnapGadget)
	// - the first post-mount task of gadget must depend on the base link
	c.Check(unlinkSnapGadget.WaitTasks(), testutil.Contains, linkTaskOfBase)
	// - "auto-connect" (MaybeRebootWaitEdge) of base must depend on "link-snap" of gadget (MaybeRebootEdge)
	c.Check(acTaskOfBase.WaitTasks(), testutil.Contains, linkTaskOfGadget)
	// - "auto-connect" (MaybeRebootWaitEdge) of gadget must depend on the last task of base (EndEdge)
	c.Check(acTaskOfGadget.WaitTasks(), testutil.Contains, lastTaskOfBase)

	// both should be transactional
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, true)

	// Since they are set up for single-reboot, the base should have restart
	// boundaries for the undo path, and gadget should have for do path.
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsGadgetKernel(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Expect restart boundaries on both
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)

	linkSnapGadget := stss[0].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapGadget, NotNil)
	mountSnapGadget := findTaskKind(stss[0].TaskSet(), "mount-snap")
	c.Assert(mountSnapGadget, NotNil)
	unlinkSnapGadget := findUnlinkTask(stss[0].TaskSet())
	c.Assert(unlinkSnapGadget, NotNil)
	linkSnapKernel := stss[1].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapKernel, NotNil)
	mountSnapKernel := findTaskKind(stss[1].TaskSet(), "mount-snap")
	c.Assert(mountSnapKernel, NotNil)
	unlinkSnapKernel := findUnlinkTask(stss[1].TaskSet())
	c.Assert(unlinkSnapKernel, NotNil)

	// Grab the tasks we need to check dependencies between
	firstTaskOfKernel := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	linkTaskOfKernel, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfKernel, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	linkTaskOfGadget, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfGadget, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	lastTaskOfGadget, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	// Things that must be correct:
	// - first local modification task of kernel must be its mount task and must depend on the gadget mount
	c.Check(firstTaskOfKernel, Equals, mountSnapKernel)
	c.Check(firstTaskOfKernel.WaitTasks(), testutil.Contains, mountSnapGadget)
	// - the gadget's remaining pre-reboot work only starts after the kernel mount phase finishes
	c.Check(unlinkSnapGadget.WaitTasks(), testutil.Contains, mountSnapKernel)
	// - the first post-mount task of kernel must depend on the gadget link
	c.Check(unlinkSnapKernel.WaitTasks(), testutil.Contains, linkTaskOfGadget)
	// - "auto-connect" (MaybeRebootWaitEdge) of gadget must depend on "link-snap" of kernel (MaybeRebootEdge)
	c.Check(acTaskOfGadget.WaitTasks(), testutil.Contains, linkTaskOfKernel)
	// - "auto-connect" (MaybeRebootWaitEdge) of kernel must depend on the last task of gadget (EndEdge)
	c.Check(acTaskOfKernel.WaitTasks(), testutil.Contains, lastTaskOfGadget)

	// both should be transactional
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, true)

	// Since they are set up for single-reboot, the gadget should have restart
	// boundaries for the undo path, and kernel should have for do path.
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsBaseGadgetKernel(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	linkSnapBase := stss[0].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapBase, NotNil)
	linkSnapGadget := stss[1].TaskSet().MaybeEdge(snapstate.MaybeRebootEdge)
	c.Assert(linkSnapGadget, NotNil)

	// Grab the tasks we need to check dependencies between
	linkTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	mountSnapBase := findTaskKind(stss[0].TaskSet(), "mount-snap")
	c.Assert(mountSnapBase, NotNil)
	unlinkSnapBase := findUnlinkTask(stss[0].TaskSet())
	c.Assert(unlinkSnapBase, NotNil)
	acTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	lastTaskOfBase, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)
	firstTaskOfGadget := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	mountSnapGadget := findTaskKind(stss[1].TaskSet(), "mount-snap")
	c.Assert(mountSnapGadget, NotNil)
	unlinkSnapGadget := findUnlinkTask(stss[1].TaskSet())
	c.Assert(unlinkSnapGadget, NotNil)
	linkTaskOfGadget, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfGadget, err := stss[1].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)
	lastTaskOfGadget, err := stss[1].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)
	firstTaskOfKernel := firstTaskAfterLocalModifications(c, stss[2].TaskSet())
	mountSnapKernel := findTaskKind(stss[2].TaskSet(), "mount-snap")
	c.Assert(mountSnapKernel, NotNil)
	unlinkSnapKernel := findUnlinkTask(stss[2].TaskSet())
	c.Assert(unlinkSnapKernel, NotNil)
	linkTaskOfKernel, err := stss[2].TaskSet().Edge(snapstate.MaybeRebootEdge)
	c.Assert(err, IsNil)
	acTaskOfKernel, err := stss[2].TaskSet().Edge(snapstate.MaybeRebootWaitEdge)
	c.Assert(err, IsNil)

	// Things that must be correct between base and gadget:
	// - first local modification task of gadget must be its mount task and must depend on the base mount
	c.Check(firstTaskOfGadget, Equals, mountSnapGadget)
	c.Check(firstTaskOfGadget.WaitTasks(), testutil.Contains, mountSnapBase)
	// - essential mounts happen before the remaining pre-reboot work starts
	c.Check(mountSnapKernel.WaitTasks(), testutil.Contains, mountSnapGadget)
	c.Check(unlinkSnapBase.WaitTasks(), testutil.Contains, mountSnapKernel)
	// - the first post-mount task of gadget must depend on the base link
	c.Check(unlinkSnapGadget.WaitTasks(), testutil.Contains, linkTaskOfBase)
	// - "auto-connect" (MaybeRebootWaitEdge) of base must depend on "link-snap" of kernel (MaybeRebootEdge)
	c.Check(acTaskOfBase.WaitTasks(), testutil.Contains, linkTaskOfKernel)
	// - "auto-connect" (MaybeRebootWaitEdge) of gadget must depend on the last task of base (EndEdge)
	c.Check(acTaskOfGadget.WaitTasks(), testutil.Contains, lastTaskOfBase)

	// Things that must be correct between gadget and kernel:
	// - first local modification task of kernel must be its mount task and must depend on the gadget mount
	c.Check(firstTaskOfKernel, Equals, mountSnapKernel)
	c.Check(firstTaskOfKernel.WaitTasks(), testutil.Contains, mountSnapGadget)
	// - the gadget's remaining pre-reboot work starts only after the base pre-reboot phase completes
	c.Check(unlinkSnapGadget.WaitTasks(), testutil.Contains, linkTaskOfBase)
	// - the first post-mount task of kernel must depend on the gadget link
	c.Check(unlinkSnapKernel.WaitTasks(), testutil.Contains, linkTaskOfGadget)
	// - "auto-connect" (MaybeRebootWaitEdge) of gadget must depend on last task of base (EndEdge)
	c.Check(acTaskOfGadget.WaitTasks(), testutil.Contains, lastTaskOfBase)
	// - "auto-connect" (MaybeRebootWaitEdge) of kernel must depend on the last task of gadget (EndEdge)
	c.Check(acTaskOfKernel.WaitTasks(), testutil.Contains, lastTaskOfGadget)

	// all three should be transactional
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet(), stss[2].TaskSet()), Equals, true)

	// Since they are set up for single-reboot, the base should have restart
	// boundaries for the undo path, and kernel should have for do path.
	c.Check(s.hasUndoRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[2].TaskSet()), Equals, true)

	// Gadget should have no boundaries
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsEarlyDownloads(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("some-app", "", snap.TypeApp),
	}
	earlyDownloads := map[string]bool{
		"core20":    true,
		"my-kernel": true,
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, earlyDownloads, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	baseLastBefore, err := stss[0].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)
	kernelLastBefore, err := stss[1].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)

	for _, sts := range stss {
		firstLocalModTask := firstTaskAfterLocalModifications(c, sts.TaskSet())
		c.Check(waitsOnTransitively(firstLocalModTask, baseLastBefore), Equals, true)
		c.Check(waitsOnTransitively(firstLocalModTask, kernelLastBefore), Equals, true)
	}
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsSnapdEarlyDownload(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snapd", "", snap.TypeSnapd),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("some-app", "", snap.TypeApp),
	}
	earlyDownloads := map[string]bool{
		"snapd": true,
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, earlyDownloads, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	snapdLastBefore, err := stss[0].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)

	baseFirstLocalMod := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	c.Check(waitsOnTransitively(baseFirstLocalMod, snapdLastBefore), Equals, true)

	appFirstLocalMod := firstTaskAfterLocalModifications(c, stss[2].TaskSet())
	c.Check(waitsOnTransitively(appFirstLocalMod, snapdLastBefore), Equals, true)

	baseLastBefore, err := stss[1].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)
	appLastBefore, err := stss[2].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)

	snapdFirstLocalMod := firstTaskAfterLocalModifications(c, stss[0].TaskSet())
	c.Check(waitsOnTransitively(snapdFirstLocalMod, baseLastBefore), Equals, false)
	c.Check(waitsOnTransitively(snapdFirstLocalMod, appLastBefore), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsSnapdAndEssentialEarlyDownloadsNoCycle(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snapd", "", snap.TypeSnapd),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("some-app", "", snap.TypeApp),
	}
	earlyDownloads := map[string]bool{
		"snapd":  true,
		"core20": true,
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, earlyDownloads, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	snapdEnd, err := stss[0].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	baseBegin, err := stss[1].TaskSet().Edge(snapstate.BeginEdge)
	c.Assert(err, IsNil)

	// if this dependency were present, it would create a cycle when snapd and
	// an essential snap are both in the early download cohort.
	c.Check(baseBegin.WaitTasks(), Not(testutil.Contains), snapdEnd)

	snapdLastBefore, err := stss[0].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)
	baseLastBefore, err := stss[1].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)

	baseFirstLocalMod := firstTaskAfterLocalModifications(c, stss[1].TaskSet())
	c.Check(waitsOnTransitively(baseFirstLocalMod, snapdLastBefore), Equals, true)

	snapdFirstLocalMod := firstTaskAfterLocalModifications(c, stss[0].TaskSet())
	c.Check(waitsOnTransitively(snapdFirstLocalMod, baseLastBefore), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsEarlyDownloadedAppWaitsAfterEssentials(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("some-app", "", snap.TypeApp),
		s.snapInstallTaskSetForSnapSetup("some-other-snap", "", snap.TypeApp),
	}
	earlyDownloads := map[string]bool{
		"core20":    true,
		"my-kernel": true,
		"some-app":  true,
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, earlyDownloads, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	finalEssential, err := stss[1].TaskSet().Edge(snapstate.EndEdge)
	c.Assert(err, IsNil)

	beginEarlyAppTask, err := stss[2].TaskSet().Edge(snapstate.BeginEdge)
	c.Assert(err, IsNil)
	c.Check(beginEarlyAppTask.WaitTasks(), Not(testutil.Contains), finalEssential)

	firstLocalModEarlyApp := firstTaskAfterLocalModifications(c, stss[2].TaskSet())
	c.Check(firstLocalModEarlyApp.WaitTasks(), testutil.Contains, finalEssential)

	beginNonEarlyAppTask, err := stss[3].TaskSet().Edge(snapstate.BeginEdge)
	c.Assert(err, IsNil)
	c.Check(beginNonEarlyAppTask.WaitTasks(), testutil.Contains, finalEssential)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsNoEarlyDownloads(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("some-app", "", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	baseLastBefore, err := stss[0].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)
	kernelLastBefore, err := stss[1].TaskSet().Edge(snapstate.LastBeforeLocalModificationsEdge)
	c.Assert(err, IsNil)

	firstLocalModTask := firstTaskAfterLocalModifications(c, stss[2].TaskSet())
	c.Check(firstLocalModTask.WaitTasks(), Not(testutil.Contains), baseLastBefore)
	c.Check(firstLocalModTask.WaitTasks(), Not(testutil.Contains), kernelLastBefore)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsSeedRefreshComponentExclusiveCandidate(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel(map[string]any{
		"base":           "core20",
		"required-snaps": []any{"some-app"},
	}))()
	observed, restore := mockSeedRefreshHooks([]string{"some-app"})
	defer restore()

	s.state.Lock()
	defer s.state.Unlock()

	tr := config.NewTransaction(s.state)
	c.Assert(tr.Set("core", "experimental.seed-refresh", true), IsNil)
	tr.Commit()

	snapsup := &snapstate.SnapSetup{
		SideInfo: &snap.SideInfo{
			RealName: "some-app",
			SnapID:   "some-app",
			Revision: snap.R(1),
		},
		Type:                        snap.TypeApp,
		ComponentExclusiveOperation: true,
	}

	downloadComp := s.state.NewTask("download-component", "...")
	setupSecurity := s.state.NewTask("setup-profiles", "...")
	setupSecurity.Set("snap-setup", snapsup)
	setupSecurity.WaitFor(downloadComp)

	componentSetupTasks := []string{downloadComp.ID()}
	setupSecurity.Set("component-setup-tasks", componentSetupTasks)

	linkComp := s.state.NewTask("link-component", "...")
	linkComp.WaitFor(setupSecurity)

	postLink := s.state.NewTask("run-hook", "...")
	postLink.WaitFor(linkComp)

	ts := state.NewTaskSet(downloadComp, setupSecurity, linkComp, postLink)
	ts.MarkEdge(downloadComp, snapstate.BeginEdge)
	ts.MarkEdge(downloadComp, snapstate.LastBeforeLocalModificationsEdge)
	ts.MarkEdge(setupSecurity, snapstate.SnapSetupEdge)
	ts.MarkEdge(postLink, snapstate.EndEdge)
	ts.JoinLane(s.state.NewLane())

	sts := snapstate.NewSnapInstallTaskSetForTest(
		snapsup,
		ts,
		[]*state.Task{downloadComp},
		nil,
		[]*state.Task{setupSecurity, linkComp},
		[]*state.Task{postLink},
	)

	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, []snapstate.SnapInstallTaskSet{sts}, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// component-exclusive operations provide only component setup tasks, so snap setup tasks must stay empty.
	observed.CheckInitialCandidates(c, snapstate.SeedRefreshCandidate{
		InstanceName:          "some-app",
		ComponentSetupTaskIDs: componentSetupTasks,
	})
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsSnapd(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snapd", "", snap.TypeSnapd),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Do not expect any restart boundaries to be set on snapd
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, false)

	// Expect them to be set on core20 as it's the boot-base
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, true)

	// Snapd should never be a part of the single-reboot transaction, we don't
	// need snapd to rollback if an issue should arise in any of the other essential snaps.
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsBootBaseAndOtherBases(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("core18", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Only the boot-base should have restart boundary.
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, false)

	// boot-base is transactional, but not with the other base and my-app
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, false)
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[2].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsForSnapWithBaseAndWithout(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snap-base", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("snap-base-app", "snap-base", snap.TypeApp),
		s.snapInstallTaskSetForSnapSetup("snap-other-app", "other-base", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// No restart boundaries
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, false)

	// no transactional lane set
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, false)
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[2].TaskSet()), Equals, false)

	// snap-base-app depends on snap-base, but snap-other-app's base
	// is not updated
	c.Check(s.setDependsOn(c, stss[1].TaskSet(), stss[0].TaskSet()), Equals, true)
	c.Check(s.setDependsOn(c, stss[2].TaskSet(), stss[0].TaskSet()), Equals, false)
	c.Check(s.setDependsOn(c, stss[2].TaskSet(), stss[1].TaskSet()), Equals, false)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsForSnapWithBootBaseAndWithout(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("snap-core20-app", "snap-core20", snap.TypeApp),
		s.snapInstallTaskSetForSnapSetup("snap-other-app", "other-base", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// Restart boundaries is set for core20 as the boot-base
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, true)
	c.Check(s.hasRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[2].TaskSet()), Equals, false)

	// Core20 is transactional, but not with the other base and app
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[1].TaskSet()), Equals, false)
	c.Check(taskSetsShareLane(stss[0].TaskSet(), stss[2].TaskSet()), Equals, false)

	// snap-core20-app depends on core20, but snap-other-app' base is
	// not updated. Yet snap-other-base still depends on core20. But there
	// is no dependency between snap-core20-app and snap-other-app
	c.Check(s.setDependsOn(c, stss[1].TaskSet(), stss[0].TaskSet()), Equals, true)  // snap-core20-app depend on core20
	c.Check(s.setDependsOn(c, stss[2].TaskSet(), stss[0].TaskSet()), Equals, true)  // snap-other-app depend on core20
	c.Check(s.setDependsOn(c, stss[2].TaskSet(), stss[1].TaskSet()), Equals, false) // snap-other-app does not depend on snap-core20-app
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsAll(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		s.snapInstallTaskSetForSnapSetup("snapd", "", snap.TypeSnapd),
		s.snapInstallTaskSetForSnapSetup("core20", "", snap.TypeBase),
		s.snapInstallTaskSetForSnapSetup("brand-gadget", "", snap.TypeGadget),
		s.snapInstallTaskSetForSnapSetup("my-kernel", "", snap.TypeKernel),
		s.snapInstallTaskSetForSnapSetup("core", "", snap.TypeOS),
		s.snapInstallTaskSetForSnapSetup("my-app", "", snap.TypeApp),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, IsNil)

	// snapd has no restart boundaries set
	c.Check(s.hasRestartBoundaries(c, stss[0].TaskSet()), Equals, false)

	// boot-base, gadget, kernel setup for single-reboot
	c.Check(s.hasDoRestartBoundaries(c, stss[1].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[1].TaskSet()), Equals, true)
	c.Check(s.hasDoRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasUndoRestartBoundaries(c, stss[2].TaskSet()), Equals, false)
	c.Check(s.hasDoRestartBoundaries(c, stss[3].TaskSet()), Equals, true)
	c.Check(s.hasUndoRestartBoundaries(c, stss[3].TaskSet()), Equals, false)

	// TypeOS (in this scenario) and TypeApp should have no restart boundaries
	c.Check(s.hasRestartBoundaries(c, stss[4].TaskSet()), Equals, false)
	c.Check(s.hasRestartBoundaries(c, stss[5].TaskSet()), Equals, false)

	// boot-base, gadget and kernel are transactional
	c.Check(taskSetsShareLane(stss[1].TaskSet(), stss[2].TaskSet(), stss[3].TaskSet()), Equals, true)
}

func (s *rebootSuite) TestArrangeSnapInstallTaskSetsFailsSplit(c *C) {
	defer snapstatetest.MockDeviceModel(MakeModel20("brand-gadget", nil))()

	s.state.Lock()
	defer s.state.Unlock()

	stss := []snapstate.SnapInstallTaskSet{
		snapstate.NewSnapInstallTaskSetForTest(nil, nil, nil, nil, nil, nil),
	}
	_, err := snapstate.ArrangeRebootAndUpdateSeed(s.state, stss, nil, snapstate.Options{DeviceCtx: s.deviceCtx(c)})
	c.Assert(err, ErrorMatches, `internal error: snap install task set has empty slices`)
}
