PR comments 5

deansheather · deansheather · commit 195fb04146ce · 2024-08-20T06:23:16.000Z
diff --git a/codersdk/workspacesdk/connector_internal_test.go b/codersdk/workspacesdk/connector_internal_test.go
@@ -226,8 +226,7 @@ func TestTailnetAPIConnector_ResumeToken(t *testing.T) {
 
 	// Fetch first token. We don't need to advance the clock since we use a
 	// channel with a single item to immediately fetch.
-	trappedTicker := newTickerTrap.MustWait(ctx)
-	trappedTicker.Release()
+	newTickerTrap.MustWait(ctx).Release()
 	// We call ticker.Reset after each token fetch to apply the refresh duration
 	// requested by the server.
 	trappedReset := tickerResetTrap.MustWait(ctx)
@@ -250,10 +249,10 @@ func TestTailnetAPIConnector_ResumeToken(t *testing.T) {
 	_ = wsConn.Close(websocket.StatusGoingAway, "test")
 
 	// Wait for the resume token to be refreshed.
-	trappedTicker = newTickerTrap.MustWait(ctx)
+	trappedTicker := newTickerTrap.MustWait(ctx)
+	// Advance the clock slightly to ensure the new JWT is different.
+	clock.Advance(time.Second).MustWait(ctx)
 	trappedTicker.Release()
-	waiter = clock.Advance(trappedTicker.Duration)
-	waiter.MustWait(ctx)
 	trappedReset = tickerResetTrap.MustWait(ctx)
 	trappedReset.Release()
 
@@ -329,8 +328,7 @@ func TestTailnetAPIConnector_ResumeTokenFailure(t *testing.T) {
 	uut.runConnector(fConn)
 
 	// Wait for the resume token to be fetched for the first time.
-	trappedTicker := newTickerTrap.MustWait(ctx)
-	trappedTicker.Release()
+	newTickerTrap.MustWait(ctx).Release()
 	trappedReset := tickerResetTrap.MustWait(ctx)
 	trappedReset.Release()
 	originalResumeToken := uut.resumeToken.Token
@@ -343,11 +341,11 @@ func TestTailnetAPIConnector_ResumeTokenFailure(t *testing.T) {
 
 	// Wait for the resume token to be refreshed, which indicates a successful
 	// reconnect.
-	trappedTicker = newTickerTrap.MustWait(ctx)
-	trappedTicker.Release()
+	trappedTicker := newTickerTrap.MustWait(ctx)
 	// Since we failed the initial reconnect and we're definitely reconnected
 	// now, the stored resume token should now be nil.
 	require.Nil(t, uut.resumeToken)
+	trappedTicker.Release()
 	trappedReset = tickerResetTrap.MustWait(ctx)
 	trappedReset.Release()
 	require.NotNil(t, uut.resumeToken)
diff --git a/tailnet/resume_test.go b/tailnet/resume_test.go
@@ -140,16 +140,14 @@ func TestResumeTokenKeyProvider(t *testing.T) {
 		t.Parallel()
 
 		id := uuid.New()
-		now := time.Now()
 		clock := quartz.NewMock(t)
-		_ = clock.Set(now)
 		provider := tailnet.NewResumeTokenKeyProvider(key, clock, tailnet.DefaultResumeTokenExpiry)
 		token, err := provider.GenerateResumeToken(id)
 		require.NoError(t, err)
 		require.NotNil(t, token)
 		require.NotEmpty(t, token.Token)
 		require.Equal(t, tailnet.DefaultResumeTokenExpiry/2, token.RefreshIn.AsDuration())
-		require.WithinDuration(t, now.Add(tailnet.DefaultResumeTokenExpiry), token.ExpiresAt.AsTime(), time.Second)
+		require.WithinDuration(t, clock.Now().Add(tailnet.DefaultResumeTokenExpiry), token.ExpiresAt.AsTime(), time.Second)
 
 		// Advance time past expiry
 		_ = clock.Advance(tailnet.DefaultResumeTokenExpiry + time.Second)
