Golang state.New函数代码示例

func (s *checkSnapSuite) TestCheckSnapGadgetMissingPrior(c *C) {
	err := os.MkdirAll(filepath.Dir(dirs.SnapFirstBootStamp), 0755)
	c.Assert(err, IsNil)
	err = ioutil.WriteFile(dirs.SnapFirstBootStamp, nil, 0644)
	c.Assert(err, IsNil)

	reset := release.MockOnClassic(false)
	defer reset()

	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	const yaml = `name: gadget
type: gadget
version: 1
`

	info, err := snap.InfoFromSnapYaml([]byte(yaml))
	c.Assert(err, IsNil)

	var openSnapFile = func(path string, si *snap.SideInfo) (*snap.Info, snap.Container, error) {
		return info, nil, nil
	}
	restore := snapstate.MockOpenSnapFile(openSnapFile)
	defer restore()

	st.Unlock()
	err = snapstate.CheckSnap(st, "snap-path", nil, 0)
	st.Lock()
	c.Check(err, ErrorMatches, "cannot find original gadget snap")
}

func (ts *taskSuite) TestProgressDefaults(c *C) {
	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	t := st.NewTask("download", "1...")

	c.Check(t.Status(), Equals, state.DoStatus)
	_, cur, tot := t.Progress()
	c.Check(cur, Equals, 0)
	c.Check(tot, Equals, 1)

	t.SetStatus(state.DoStatus)
	_, cur, tot = t.Progress()
	c.Check(cur, Equals, 0)
	c.Check(tot, Equals, 1)

	t.SetStatus(state.DoneStatus)
	_, cur, tot = t.Progress()
	c.Check(cur, Equals, 1)
	c.Check(tot, Equals, 1)

	t.SetStatus(state.ErrorStatus)
	_, cur, tot = t.Progress()
	c.Check(cur, Equals, 1)
	c.Check(tot, Equals, 1)
}

func (s *deviceMgrSuite) SetUpTest(c *C) {
	dirs.SetRootDir(c.MkDir())

	rootPrivKey, _ := assertstest.GenerateKey(1024)
	storePrivKey, _ := assertstest.GenerateKey(752)
	s.storeSigning = assertstest.NewStoreStack("canonical", rootPrivKey, storePrivKey)
	s.state = state.New(nil)

	db, err := asserts.OpenDatabase(&asserts.DatabaseConfig{
		Backstore: asserts.NewMemoryBackstore(),
		Trusted:   s.storeSigning.Trusted,
	})
	c.Assert(err, IsNil)

	s.state.Lock()
	assertstate.ReplaceDB(s.state, db)
	s.state.Unlock()

	err = db.Add(s.storeSigning.StoreAccountKey(""))
	c.Assert(err, IsNil)

	mgr, err := devicestate.Manager(s.state)
	c.Assert(err, IsNil)

	s.db = db
	s.mgr = mgr

	s.state.Lock()
	snapstate.ReplaceStore(s.state, &fakeStore{
		state: s.state,
		db:    s.storeSigning,
	})
	s.state.Unlock()
}

func (ts *taskRunnerSuite) TestExternalAbort(c *C) {
	sb := &stateBackend{}
	st := state.New(sb)
	r := state.NewTaskRunner(st)
	defer r.Stop()

	ch := make(chan bool)
	r.AddHandler("blocking", func(t *state.Task, tb *tomb.Tomb) error {
		ch <- true
		<-tb.Dying()
		return nil
	}, nil)

	st.Lock()
	chg := st.NewChange("install", "...")
	t := st.NewTask("blocking", "...")
	chg.AddTask(t)
	st.Unlock()

	r.Ensure()
	<-ch

	st.Lock()
	chg.Abort()
	st.Unlock()

	// The Abort above must make Ensure kill the task, or this will never end.
	ensureChange(c, r, sb, chg)
}

func (s *checkSnapSuite) TestCheckSnapGadgetCannotBeInstalledOnClassic(c *C) {
	reset := release.MockOnClassic(true)
	defer reset()

	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	const yaml = `name: gadget
type: gadget
version: 1
`

	info, err := snap.InfoFromSnapYaml([]byte(yaml))
	c.Assert(err, IsNil)

	var openSnapFile = func(path string, si *snap.SideInfo) (*snap.Info, snap.Container, error) {
		return info, nil, nil
	}
	restore := snapstate.MockOpenSnapFile(openSnapFile)
	defer restore()

	st.Unlock()
	err = snapstate.CheckSnap(st, "snap-path", nil, nil, snapstate.Flags{})
	st.Lock()
	c.Check(err, ErrorMatches, "cannot install a gadget snap on classic")
}

func loadState(backend state.Backend) (*state.State, error) {
	if !osutil.FileExists(dirs.SnapStateFile) {
		// fail fast, mostly interesting for tests, this dir is setup
		// by the snapd package
		stateDir := filepath.Dir(dirs.SnapStateFile)
		if !osutil.IsDirectory(stateDir) {
			return nil, fmt.Errorf("fatal: directory %q must be present", stateDir)
		}
		s := state.New(backend)
		patch.Init(s)
		return s, nil
	}

	r, err := os.Open(dirs.SnapStateFile)
	if err != nil {
		return nil, fmt.Errorf("cannot read the state file: %s", err)
	}
	defer r.Close()

	s, err := state.ReadState(backend, r)
	if err != nil {
		return nil, err
	}

	// one-shot migrations
	err = patch.Apply(s)
	if err != nil {
		return nil, err
	}
	return s, nil
}

func (s *assertMgrSuite) SetUpTest(c *C) {
	dirs.SetRootDir(c.MkDir())

	rootPrivKey, _ := assertstest.GenerateKey(1024)
	storePrivKey, _ := assertstest.GenerateKey(752)
	s.storeSigning = assertstest.NewStoreStack("can0nical", rootPrivKey, storePrivKey)
	s.restore = sysdb.InjectTrusted(s.storeSigning.Trusted)

	dev1PrivKey, _ := assertstest.GenerateKey(752)
	s.dev1Acct = assertstest.NewAccount(s.storeSigning, "developer1", nil, "")
	err := s.storeSigning.Add(s.dev1Acct)
	c.Assert(err, IsNil)

	// developer signing
	dev1AcctKey := assertstest.NewAccountKey(s.storeSigning, s.dev1Acct, nil, dev1PrivKey.PublicKey(), "")
	err = s.storeSigning.Add(dev1AcctKey)
	c.Assert(err, IsNil)

	s.dev1Signing = assertstest.NewSigningDB(s.dev1Acct.AccountID(), dev1PrivKey)

	s.state = state.New(nil)
	mgr, err := assertstate.Manager(s.state)
	c.Assert(err, IsNil)
	s.mgr = mgr

	s.state.Lock()
	snapstate.ReplaceStore(s.state, &fakeStore{
		state: s.state,
		db:    s.storeSigning,
	})
	s.state.Unlock()
}

func (s *repositorySuite) TestAddHandlerGenerator(c *C) {
	repository := newRepository()

	var calledContext *Context
	mockHandlerGenerator := func(context *Context) Handler {
		calledContext = context
		return hooktest.NewMockHandler()
	}

	// Verify that a handler generator can be added to the repository
	repository.addHandlerGenerator(regexp.MustCompile("test-hook"), mockHandlerGenerator)

	state := state.New(nil)
	state.Lock()
	task := state.NewTask("test-task", "my test task")
	setup := &HookSetup{Snap: "test-snap", Revision: snap.R(1), Hook: "test-hook"}
	context := &Context{task: task, setup: setup}
	state.Unlock()

	c.Assert(context, NotNil)

	// Verify that the handler can be generated
	handlers := repository.generateHandlers(context)
	c.Check(handlers, HasLen, 1)
	c.Check(calledContext, DeepEquals, context)

	// Add another handler
	repository.addHandlerGenerator(regexp.MustCompile(".*-hook"), mockHandlerGenerator)

	// Verify that two handlers are generated for the test-hook, now
	handlers = repository.generateHandlers(context)
	c.Check(handlers, HasLen, 2)
	c.Check(calledContext, DeepEquals, context)
}

func (cs *changeSuite) TestCloseReadyWhenTasksReady(c *C) {
	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	chg := st.NewChange("install", "...")
	t1 := st.NewTask("download", "...")
	t2 := st.NewTask("download", "...")
	chg.AddTask(t1)
	chg.AddTask(t2)

	select {
	case <-chg.Ready():
		c.Fatalf("Change should not be ready")
	default:
	}

	t1.SetStatus(state.DoneStatus)

	select {
	case <-chg.Ready():
		c.Fatalf("Change should not be ready")
	default:
	}

	t2.SetStatus(state.DoneStatus)

	select {
	case <-chg.Ready():
	default:
		c.Fatalf("Change should be ready")
	}
}

func (ss *stateSuite) TestPrune(c *C) {
	st := state.New(&fakeStateBackend{})
	st.Lock()
	defer st.Unlock()

	now := time.Now()
	pruneWait := 1 * time.Hour
	abortWait := 3 * time.Hour

	unset := time.Time{}

	t1 := st.NewTask("foo", "...")
	t2 := st.NewTask("foo", "...")
	t3 := st.NewTask("foo", "...")
	t4 := st.NewTask("foo", "...")

	chg1 := st.NewChange("abort", "...")
	chg1.AddTask(t1)
	state.MockChangeTimes(chg1, now.Add(-abortWait), unset)

	chg2 := st.NewChange("prune", "...")
	chg2.AddTask(t2)
	c.Assert(chg2.Status(), Equals, state.DoStatus)
	state.MockChangeTimes(chg2, now.Add(-pruneWait), now.Add(-pruneWait))

	chg3 := st.NewChange("ready-but-recent", "...")
	chg3.AddTask(t3)
	state.MockChangeTimes(chg3, now.Add(-pruneWait), now.Add(-pruneWait/2))

	chg4 := st.NewChange("old-but-not-ready", "...")
	chg4.AddTask(t4)
	state.MockChangeTimes(chg4, now.Add(-pruneWait/2), unset)

	// unlinked task
	t5 := st.NewTask("unliked", "...")
	c.Check(st.Task(t5.ID()), IsNil)
	state.MockTaskTimes(t5, now.Add(-pruneWait), now.Add(-pruneWait))

	st.Prune(pruneWait, abortWait)

	c.Assert(st.Change(chg1.ID()), Equals, chg1)
	c.Assert(st.Change(chg2.ID()), IsNil)
	c.Assert(st.Change(chg3.ID()), Equals, chg3)
	c.Assert(st.Change(chg4.ID()), Equals, chg4)

	c.Assert(st.Task(t1.ID()), Equals, t1)
	c.Assert(st.Task(t2.ID()), IsNil)
	c.Assert(st.Task(t3.ID()), Equals, t3)
	c.Assert(st.Task(t4.ID()), Equals, t4)

	c.Assert(chg1.Status(), Equals, state.HoldStatus)
	c.Assert(chg3.Status(), Equals, state.DoStatus)
	c.Assert(chg4.Status(), Equals, state.DoStatus)

	c.Assert(t1.Status(), Equals, state.HoldStatus)
	c.Assert(t3.Status(), Equals, state.DoStatus)
	c.Assert(t4.Status(), Equals, state.DoStatus)

	c.Check(st.TaskCount(), Equals, 3)
}

func (ts *taskRunnerSuite) TestStopAskForRetry(c *C) {
	sb := &stateBackend{}
	st := state.New(sb)
	r := state.NewTaskRunner(st)
	defer r.Stop()

	ch := make(chan bool)
	r.AddHandler("ask-for-retry", func(t *state.Task, tb *tomb.Tomb) error {
		ch <- true
		<-tb.Dying()
		// ask for retry
		return &state.Retry{}
	}, nil)

	st.Lock()
	chg := st.NewChange("install", "...")
	t := st.NewTask("ask-for-retry", "...")
	chg.AddTask(t)
	st.Unlock()

	r.Ensure()
	<-ch
	r.Stop()

	st.Lock()
	defer st.Unlock()
	c.Check(t.Status(), Equals, state.DoingStatus)
}

func (ss *stateSuite) TestNewTaskAndTasks(c *C) {
	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	chg1 := st.NewChange("install", "...")
	t11 := st.NewTask("check", "...")
	chg1.AddTask(t11)
	t12 := st.NewTask("inst", "...")
	chg1.AddTask(t12)

	chg2 := st.NewChange("remove", "...")
	t21 := st.NewTask("check", "...")
	t22 := st.NewTask("rm", "...")
	chg2.AddTask(t21)
	chg2.AddTask(t22)

	tasks := st.Tasks()
	c.Check(tasks, HasLen, 4)

	expected := map[string]*state.Task{
		t11.ID(): t11,
		t12.ID(): t12,
		t21.ID(): t21,
		t22.ID(): t22,
	}

	for _, t := range tasks {
		c.Check(t, Equals, expected[t.ID()])
	}
}

func (ss *stateSuite) TestTaskNoTask(c *C) {
	st := state.New(nil)
	st.Lock()
	defer st.Unlock()

	c.Check(st.Task("1"), IsNil)
}

func (ss *stateSuite) TestCheckpointPreserveLastIds(c *C) {
	b := new(fakeStateBackend)
	st := state.New(b)
	st.Lock()

	st.NewChange("install", "...")
	st.NewTask("download", "...")
	st.NewTask("download", "...")

	c.Assert(st.NewLane(), Equals, 1)

	// implicit checkpoint
	st.Unlock()

	c.Assert(b.checkpoints, HasLen, 1)

	buf := bytes.NewBuffer(b.checkpoints[0])

	st2, err := state.ReadState(nil, buf)
	c.Assert(err, IsNil)

	st2.Lock()
	defer st2.Unlock()

	c.Assert(st2.NewTask("download", "...").ID(), Equals, "3")
	c.Assert(st2.NewChange("install", "...").ID(), Equals, "2")

	c.Assert(st2.NewLane(), Equals, 2)

}

func (ss *stateSuite) TestCheckpointPreserveCleanStatus(c *C) {
	b := new(fakeStateBackend)
	st := state.New(b)
	st.Lock()

	chg := st.NewChange("install", "...")
	t := st.NewTask("download", "...")
	chg.AddTask(t)
	t.SetStatus(state.DoneStatus)
	t.SetClean()

	// implicit checkpoint
	st.Unlock()

	c.Assert(b.checkpoints, HasLen, 1)

	buf := bytes.NewBuffer(b.checkpoints[0])

	st2, err := state.ReadState(nil, buf)
	c.Assert(err, IsNil)

	st2.Lock()
	defer st2.Unlock()

	chg2 := st2.Change(chg.ID())
	t2 := st2.Task(t.ID())

	c.Assert(chg2.IsClean(), Equals, true)
	c.Assert(t2.IsClean(), Equals, true)
}