notifier: refactor, add more unit tests

johnstcn · johnstcn · commit 8313a52d345a · 2022-05-13T09:34:09.000Z
diff --git a/coderd/autobuild/notify/notifier.go b/coderd/autobuild/notify/notifier.go
@@ -64,10 +64,15 @@ func (n *Notifier) pollOnce(tick time.Time) {
 
 	timeRemaining := deadline.Sub(tick)
 	for _, tock := range n.countdown {
-		if timeRemaining <= tock && !n.notifiedAt[tock] {
-			callback()
-			n.notifiedAt[tock] = true
+		if n.notifiedAt[tock] {
+			continue
 		}
+		if timeRemaining > tock {
+			continue
+		}
+		callback()
+		n.notifiedAt[tock] = true
+		return
 	}
 }
 
diff --git a/coderd/autobuild/notify/notifier_test.go b/coderd/autobuild/notify/notifier_test.go
@@ -27,43 +27,51 @@ func TestNotifier(t *testing.T) {
 		{
 			Name:              "zero deadline",
 			Countdown:         durations(),
-			Ticks:             ticks(now, 0),
+			Ticks:             fakeTicker(now, time.Second, 0),
 			ConditionDeadline: time.Time{},
 			NumConditions:     1,
 			NumCallbacks:      0,
 		},
 		{
 			Name:              "no calls",
 			Countdown:         durations(),
-			Ticks:             ticks(now, 0),
+			Ticks:             fakeTicker(now, time.Second, 0),
 			ConditionDeadline: now,
 			NumConditions:     1,
 			NumCallbacks:      0,
 		},
 		{
 			Name:              "exactly one call",
 			Countdown:         durations(time.Second),
-			Ticks:             ticks(now, 2),
-			ConditionDeadline: now.Add(2 * time.Second),
+			Ticks:             fakeTicker(now, time.Second, 1),
+			ConditionDeadline: now.Add(time.Second),
 			NumConditions:     2,
 			NumCallbacks:      1,
 		},
 		{
 			Name:              "two calls",
 			Countdown:         durations(4*time.Second, 2*time.Second),
-			Ticks:             ticks(now, 5),
+			Ticks:             fakeTicker(now, time.Second, 5),
 			ConditionDeadline: now.Add(5 * time.Second),
 			NumConditions:     6,
 			NumCallbacks:      2,
 		},
 		{
 			Name:              "wrong order should not matter",
 			Countdown:         durations(2*time.Second, 4*time.Second),
-			Ticks:             ticks(now, 5),
+			Ticks:             fakeTicker(now, time.Second, 5),
 			ConditionDeadline: now.Add(5 * time.Second),
 			NumConditions:     6,
 			NumCallbacks:      2,
 		},
+		{
+			Name:              "ssh autostop notify",
+			Countdown:         durations(5*time.Minute, time.Minute),
+			Ticks:             fakeTicker(now, 30*time.Second, 120),
+			ConditionDeadline: now.Add(30 * time.Minute),
+			NumConditions:     121,
+			NumCallbacks:      2,
+		},
 	}
 
 	for _, testCase := range testCases {
@@ -91,6 +99,7 @@ func TestNotifier(t *testing.T) {
 			close(ch)
 			wg.Wait()
 			require.Equal(t, testCase.NumCallbacks, numCalls.Load())
+			require.Equal(t, testCase.NumConditions, numConditions.Load())
 		})
 	}
 }
@@ -99,11 +108,10 @@ func durations(ds ...time.Duration) []time.Duration {
 	return ds
 }
 
-func ticks(t time.Time, n int) []time.Time {
-	ts := make([]time.Time, n+1)
-	ts = append(ts, t)
-	for i := 0; i < n; i++ {
-		ts = append(ts, t.Add(time.Duration(n)*time.Second))
+func fakeTicker(t time.Time, d time.Duration, n int) []time.Time {
+	ts := make([]time.Time, 0)
+	for i := 1; i <= n; i++ {
+		ts = append(ts, t.Add(time.Duration(n)*d))
 	}
 	return ts
 }

Original file line number	Diff line number	Diff line change
`@@ -64,10 +64,15 @@ func (n *Notifier) pollOnce(tick time.Time) {`
`64`	`64`
`65`	`65`	`timeRemaining := deadline.Sub(tick)`
`66`	`66`	`for _, tock := range n.countdown {`
`67`		`- if timeRemaining <= tock && !n.notifiedAt[tock] {`
`68`		`- callback()`
`69`		`- n.notifiedAt[tock] = true`
	`67`	`+ if n.notifiedAt[tock] {`
	`68`	`+ continue`
`70`	`69`	`}`
	`70`	`+ if timeRemaining > tock {`
	`71`	`+ continue`
	`72`	`+ }`
	`73`	`+ callback()`
	`74`	`+ n.notifiedAt[tock] = true`
	`75`	`+ return`
`71`	`76`	`}`
`72`	`77`	`}`
`73`	`78`