updatehub: fix state change transition

UpdateHub · Jul 26, 2018 · cb410c1 · cb410c1
1 parent d0dbbe0
commit cb410c1
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 33 deletions.
diff --git a/updatehub/daemon.go b/updatehub/daemon.go
@@ -27,6 +27,8 @@ func (d *Daemon) Run() int {
 	for {
 		nextState := d.uh.ProcessCurrentState()
 
+		d.uh.SetState(nextState)
+
 		if d.stop || nextState.ID() == UpdateHubStateExit {
 			if finalState, _ := nextState.(*ExitState); finalState != nil {
 				return finalState.exitCode

diff --git a/updatehub/updatehub.go b/updatehub/updatehub.go
@@ -231,48 +231,42 @@ func (uh *UpdateHub) ProcessCurrentState() State {
 		}
 
 		state, _ := uh.state.Handle(uh)
-		uh.state = state
-	} else {
-		flow, err := uh.stateChangeCallback(uh.state, "enter")
+		return state
+	}
 
-		uh.ReportCurrentState()
-		// this must be done after the report, because the report uses it
-		uh.previousState = uh.state
+	flow, err := uh.stateChangeCallback(uh.state, "enter")
 
-		if err != nil {
-			log.Error(err)
-			uh.state = NewErrorState(uh.state.ApiClient(), nil, NewTransientError(err))
-			return uh.state
-		}
+	uh.ReportCurrentState()
+	// this must be done after the report, because the report uses it
+	uh.previousState = uh.state
 
-		if flow != nil {
-			uh.state = flow
-			return uh.state
-		}
+	if err != nil {
+		log.Error(err)
+		return NewErrorState(uh.state.ApiClient(), nil, NewTransientError(err))
+	}
 
-		state, cancel := uh.state.Handle(uh)
+	if flow != nil {
+		return flow
+	}
 
-		flow, err = uh.stateChangeCallback(uh.state, "leave")
-		if err != nil {
-			log.Error(err)
-			uh.state = NewErrorState(uh.state.ApiClient(), nil, NewTransientError(err))
-			return uh.state
-		}
+	state, cancel := uh.state.Handle(uh)
 
-		if flow != nil {
-			uh.state = flow
-			return uh.state
-		}
+	flow, err = uh.stateChangeCallback(uh.state, "leave")
+	if err != nil {
+		log.Error(err)
+		return NewErrorState(uh.state.ApiClient(), nil, NewTransientError(err))
+	}
 
-		cs, ok := uh.state.(*CancellableState)
-		if cancel && ok {
-			uh.state = cs.NextState()
-		} else {
-			uh.state = state
-		}
+	if flow != nil {
+		return flow
 	}
 
-	return uh.state
+	cs, ok := state.(*CancellableState)
+	if cancel && ok {
+		return cs.NextState()
+	}
+
+	return state
 }
 
 type Controller interface {

diff --git a/updatehub/updatehub_test.go b/updatehub/updatehub_test.go
@@ -371,6 +371,45 @@ func TestProcessCurrentStateWithLeaveFlow(t *testing.T) {
 	cm.AssertExpectations(t)
 }
 
+func TestProcessCurrentStateWithCancellable(t *testing.T) {
+	aim := &activeinactivemock.ActiveInactiveMock{}
+	rm := &reportermock.ReporterMock{}
+	cm := &cmdlinemock.CmdLineExecuterMock{}
+	fs := afero.NewMemMapFs()
+
+	afero.WriteFile(fs, "/usr/share/updatehub/state-change-callback", []byte("a"), 0755)
+
+	expectedCallOrder := []string{"enter", "leave"}
+	callOrder := []string{}
+
+	cm.On("Execute", "/usr/share/updatehub/state-change-callback enter poll").Run(func(args mock.Arguments) {
+		callOrder = append(callOrder, "enter")
+	}).Return([]byte(""), nil).Once()
+	cm.On("Execute", "/usr/share/updatehub/state-change-callback leave poll").Run(func(args mock.Arguments) {
+		callOrder = append(callOrder, "leave")
+	}).Return([]byte(""), nil).Once()
+
+	poll := NewPollState(0)
+
+	uh, _ := newTestUpdateHub(poll, aim)
+	uh.CmdLineExecuter = cm
+	uh.Store = fs
+	uh.Reporter = rm
+
+	go func() {
+		assert.True(t, poll.Cancel(true, NewIdleState()))
+	}()
+
+	nextState := uh.ProcessCurrentState()
+
+	assert.IsType(t, &IdleState{}, nextState)
+	assert.Equal(t, expectedCallOrder, callOrder)
+
+	aim.AssertExpectations(t)
+	cm.AssertExpectations(t)
+	cm.AssertExpectations(t)
+}
+
 func TestNewUpdateHub(t *testing.T) {
 	gitversion := "2.1"
 	memFs := afero.NewMemMapFs()