Merge pull request #186 from gustavosbarreto/fix_probe_asap

Fix probe ASAP
UpdateHub · Jan 31, 2018 · 6faa852 · 6faa852
2 parents ec581e0 + e9648b1
commit 6faa852
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 4 deletions.
diff --git a/updatehub/probe_state.go b/updatehub/probe_state.go
@@ -66,6 +66,10 @@ func (state *ProbeState) Handle(uh *UpdateHub) (State, bool) {
 		uh.Settings.PollingRetries = 0
 	}
 
+	if uh.Settings.ProbeASAP {
+		uh.Settings.ProbeASAP = false
+	}
+
 	uh.Settings.LastPoll = time.Now()
 	uh.Settings.ExtraPollingInterval = 0
 	uh.Settings.Save(uh.Store)

diff --git a/updatehub/probe_state_test.go b/updatehub/probe_state_test.go
@@ -286,6 +286,55 @@ func TestStateProbeNotHavePendingDownload(t *testing.T) {
 	cm.AssertExpectations(t)
 }
 
+func TestStateProbeASAP(t *testing.T) {
+	aim := &activeinactivemock.ActiveInactiveMock{}
+	om := &objectmock.ObjectMock{}
+	cm := &controllermock.ControllerMock{}
+
+	mode := installmodes.RegisterInstallMode(installmodes.InstallMode{
+		Name:              "test",
+		CheckRequirements: func() error { return nil },
+		GetObject:         func() interface{} { return om },
+	})
+	defer mode.Unregister()
+
+	um, err := metadata.NewUpdateMetadata([]byte(validJSONMetadata))
+	assert.NoError(t, err)
+
+	apiClient := client.NewApiClient("address")
+
+	uh, err := newTestUpdateHub(NewProbeState(apiClient), aim)
+	assert.NoError(t, err)
+
+	uh.Controller = cm
+
+	uh.Store.Remove(uh.Settings.RuntimeSettingsPath)
+
+	uh.Settings.ProbeASAP = true
+
+	sha256sum := sha256.Sum256([]byte(validJSONMetadata))
+	signature, _ := rsa.SignPKCS1v15(rand.Reader, testPrivateKey, crypto.SHA256, sha256sum[:])
+
+	cm.On("ProbeUpdate", apiClient, 0).Return(um, signature, time.Duration(0), nil)
+
+	next, _ := uh.GetState().Handle(uh)
+
+	assert.Equal(t, NewDownloadingState(apiClient, um, &ProgressTrackerImpl{}), next)
+
+	data, err := afero.ReadFile(uh.Store, uh.Settings.RuntimeSettingsPath)
+	assert.NoError(t, err)
+	assert.True(t, strings.Contains(string(data), "ProbeASAP=false"))
+	assert.True(t, strings.Contains(string(data), "Retries=0"))
+	assert.True(t, strings.Contains(string(data), "ExtraInterval=0"))
+	// timestamps are relative to "Now()" so just test if they were written
+	assert.True(t, strings.Contains(string(data), "FirstPoll="))
+	assert.True(t, strings.Contains(string(data), "LastPoll="))
+
+	aim.AssertExpectations(t)
+	om.AssertExpectations(t)
+	cm.AssertExpectations(t)
+}
+
 func TestStateProbeToMap(t *testing.T) {
 	state := NewProbeState(client.NewApiClient("address"))
 

diff --git a/updatehub/updatehub.go b/updatehub/updatehub.go
@@ -624,8 +624,6 @@ func (uh *UpdateHub) Start() {
 
 	if uh.Settings.ProbeASAP {
 		uh.state = NewProbeState(uh.DefaultApiClient)
-		uh.Settings.ProbeASAP = false
-		uh.Settings.Save(uh.Store)
 	}
 }
 

diff --git a/updatehub/updatehub_test.go b/updatehub/updatehub_test.go
@@ -1983,11 +1983,11 @@ func TestStart(t *testing.T) {
 			(time.Time{}).UTC(),
 			&ProbeState{},
 			func(t *testing.T, uh *UpdateHub, state State) {
-				assert.Equal(t, false, uh.Settings.ProbeASAP)
+				assert.Equal(t, true, uh.Settings.ProbeASAP)
 
 				data, err := afero.ReadFile(uh.Store, uh.Settings.RuntimeSettingsPath)
 				assert.NoError(t, err)
-				assert.True(t, strings.Contains(string(data), "ProbeASAP=false"))
+				assert.True(t, strings.Contains(string(data), "ProbeASAP=true"))
 				assert.True(t, strings.Contains(string(data), "Retries=0"))
 				assert.True(t, strings.Contains(string(data), "ExtraInterval=0"))
 				// timestamps are relative to "Now()" so just test if they were written