From 1b6b22ef4f1f7e542cd418889446770aade3cf6a Mon Sep 17 00:00:00 2001
From: Johan Brandhorst-Satzkorn <johan.brandhorst@gmail.com>
Date: Fri, 20 Dec 2024 12:02:22 -0800
Subject: [PATCH] internal/daemon/worker: improve runtime and reduce noise
 (#5387)

TestRotationTicking would previously wait for 5 rotations, which is
pretty slow considering each rotation takes 20 seconds. Reduce
the number to 3, and use require instead of assert to remove
noise during flaky test runs.

This doesn't fix the flakiness of this test, but it should reduce the log spam.
---
 internal/daemon/worker/auth_rotation_test.go | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/internal/daemon/worker/auth_rotation_test.go b/internal/daemon/worker/auth_rotation_test.go
index e6dc4f0f2..a8a410cf7 100644
--- a/internal/daemon/worker/auth_rotation_test.go
+++ b/internal/daemon/worker/auth_rotation_test.go
@@ -28,6 +28,8 @@ import (
 // TestRotationTicking ensures that we see new credential information for a
 // worker on both the controller side and worker side in a periodic fashion
 func TestRotationTicking(t *testing.T) {
+	t.Parallel()
+
 	require, assert := require.New(t), assert.New(t)
 	logger := hclog.New(&hclog.LoggerOptions{
 		Level: hclog.Trace,
@@ -95,7 +97,7 @@ func TestRotationTicking(t *testing.T) {
 	// Verify we see authorized credentials now
 	auths, err = workerAuthRepo.List(c.Context(), (*types.NodeInformation)(nil))
 	require.NoError(err)
-	assert.Len(auths, 2)
+	require.Len(auths, 2)
 	// Fetch creds and store current key
 	currNodeCreds, err := types.LoadNodeCredentials(
 		w.Context(),
@@ -112,21 +114,29 @@ func TestRotationTicking(t *testing.T) {
 	// Make sure we have a failsafe in case the below loop never finds what it's
 	// looking for; at expiration the List will fail and we'll hit the Require,
 	// exiting
-	deadlineCtx, deadlineCtxCancel := context.WithTimeout(c.Context(), 10*time.Minute)
+	testTimeout := 10 * time.Minute
+	testDeadline, ok := t.Deadline()
+	if ok && time.Until(testDeadline) < testTimeout {
+		// If the test deadline is less than the default timeout, use the deadline
+		// Give a minute buffer for shutdowns and other test cleanup
+		testTimeout = time.Until(testDeadline) - time.Minute
+	}
+	deadlineCtx, deadlineCtxCancel := context.WithTimeout(c.Context(), testTimeout)
 	defer deadlineCtxCancel()
 
 	rotationCount := 2
 	for {
-		if rotationCount > 5 {
+		if rotationCount > 3 {
 			break
 		}
 		nextRotation := w.Worker().AuthRotationNextRotation.Load()
-		time.Sleep((*nextRotation).Sub(time.Now()) + 5*time.Second)
+		// Sleep until 5 seconds after next rotation
+		time.Sleep(time.Until(*nextRotation) + 5*time.Second)
 
 		// Verify we see 2- after credentials have rotated, we should see current and previous
 		auths, err = workerAuthRepo.List(deadlineCtx, (*types.NodeInformation)(nil))
 		require.NoError(err)
-		assert.Len(auths, 2)
+		require.Len(auths, 2)
 
 		// Fetch creds and compare current key
 		currNodeCreds, err := types.LoadNodeCredentials(