Optionally cancel context in TestWaitForWatcher to test timeouts (#4425)

Author: pchila

internal/pkg/agent/application/upgrade/upgrade_test.go

@@ -748,27 +748,25 @@ func TestWaitForWatcher(t *testing.T) {
 	wantErrWatcherNotStarted := func(t assert.TestingT, err error, i ...interface{}) bool {
 		return assert.ErrorIs(t, err, ErrWatcherNotStarted, i)
 	}
+
 	tests := []struct {
-		name                     string
-		states                   []details.State
-		stateChangeInterval      time.Duration
-		timeout                  time.Duration
-		expirationAfterLastState time.Duration
-		wantErr                  assert.ErrorAssertionFunc
+		name                string
+		states              []details.State
+		stateChangeInterval time.Duration
+		cancelWaitContext   bool
+		wantErr             assert.ErrorAssertionFunc
 	}{
 		{
-			name:                     "Happy path: watcher is watching already",
-			states:                   []details.State{details.StateWatching},
-			stateChangeInterval:      1 * time.Millisecond,
-			timeout:                  500 * time.Millisecond,
-			expirationAfterLastState: 100 * time.Millisecond,
-			wantErr:                  assert.NoError,
+			name:                "Happy path: watcher is watching already",
+			states:              []details.State{details.StateWatching},
+			stateChangeInterval: 1 * time.Millisecond,
+			wantErr:             assert.NoError,
 		},
 		{
 			name:                "Sad path: watcher is never starting",
 			states:              []details.State{details.StateReplacing},
 			stateChangeInterval: 1 * time.Millisecond,
-			timeout:             50 * time.Millisecond,
+			cancelWaitContext:   true,
 			wantErr:             wantErrWatcherNotStarted,
 		},
 		{
@@ -782,16 +780,14 @@ func TestWaitForWatcher(t *testing.T) {
 				details.StateRestarting,
 				details.StateWatching,
 			},
-			stateChangeInterval:      1 * time.Millisecond,
-			timeout:                  500 * time.Millisecond,
-			expirationAfterLastState: 10 * time.Millisecond,
-			wantErr:                  assert.NoError,
+			stateChangeInterval: 1 * time.Millisecond,
+			wantErr:             assert.NoError,
 		},
 		{
 			name:                "Timeout: marker is never created",
 			states:              nil,
 			stateChangeInterval: 1 * time.Millisecond,
-			timeout:             50 * time.Millisecond,
+			cancelWaitContext:   true,
 			wantErr:             wantErrWatcherNotStarted,
 		},
 		{
@@ -805,8 +801,8 @@ func TestWaitForWatcher(t *testing.T) {
 				details.StateRestarting,
 			},
 
-			stateChangeInterval: 5 * time.Millisecond,
-			timeout:             20 * time.Millisecond,
+			stateChangeInterval: 1 * time.Millisecond,
+			cancelWaitContext:   true,
 			wantErr:             wantErrWatcherNotStarted,
 		},
 	}
@@ -817,18 +813,22 @@ func TestWaitForWatcher(t *testing.T) {
 			if !ok {
 				deadline = time.Now().Add(5 * time.Second)
 			}
-			testCtx, cancel := context.WithDeadline(context.TODO(), deadline)
-			defer cancel()
+			testCtx, testCancel := context.WithDeadline(context.Background(), deadline)
+			defer testCancel()
 
 			tmpDir := t.TempDir()
 			updMarkerFilePath := filepath.Join(tmpDir, markerFilename)
 
 			waitContext, waitCancel := context.WithCancel(testCtx)
+			defer waitCancel()
+
+			fakeTimeout := 30 * time.Second
+
 			// in order to take timing out of the equation provide a context that we can cancel manually
 			// still assert that the parent context and timeout passed are correct
 			var createContextFunc createContextWithTimeout = func(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
 				assert.Same(t, testCtx, ctx, "parent context should be the same as the waitForWatcherCall")
-				assert.Equal(t, tt.timeout, timeout, "timeout used in new context should be the same as testcase")
+				assert.Equal(t, fakeTimeout, timeout, "timeout used in new context should be the same as testcase")
 
 				return waitContext, waitCancel
 			}
@@ -848,6 +848,8 @@ func TestWaitForWatcher(t *testing.T) {
 
 			wg.Add(1)
 
+			// worker goroutine: writes out additional states while the test is blocked on waitOnWatcher() call and expires
+			// the wait context if cancelWaitContext is set to true. Timing of the goroutine is driven by stateChangeInterval.
 			go func() {
 				defer wg.Done()
 				tick := time.NewTicker(tt.stateChangeInterval)
@@ -860,16 +862,15 @@ func TestWaitForWatcher(t *testing.T) {
 						writeState(t, updMarkerFilePath, state)
 					}
 				}
-				<-time.After(tt.expirationAfterLastState)
-				waitCancel()
+				if tt.cancelWaitContext {
+					<-tick.C
+					waitCancel()
+				}
 			}()
 
 			log, _ := logger.NewTesting(tt.name)
 
-			tt.wantErr(t, waitForWatcherWithTimeoutCreationFunc(testCtx, log, updMarkerFilePath, tt.timeout, createContextFunc), fmt.Sprintf("waitForWatcher %s, %v, %s, %s)", updMarkerFilePath, tt.states, tt.stateChangeInterval, tt.timeout))
-
-			// cancel context
-			cancel()
+			tt.wantErr(t, waitForWatcherWithTimeoutCreationFunc(testCtx, log, updMarkerFilePath, fakeTimeout, createContextFunc), fmt.Sprintf("waitForWatcher %s, %v, %s, %s)", updMarkerFilePath, tt.states, tt.stateChangeInterval, fakeTimeout))
 
 			// wait for goroutines to finish
 			wg.Wait()