From 92f7d8cfb20a3254d4c8e3cd478f78dd4e614893 Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Mon, 6 Nov 2023 16:03:25 -0500
Subject: [PATCH] Minor test fixes (#4004)

* Minor test fixes

The auth rotation test fix allows us to keep checking if we did not yet
see the expected bytes in case the call to lookup succeeds but we are
late in rotating due to parallelism.

The api credentials test fix allows us to check if at _least_ one
rotation has happened. That way if two rotations happen before we check
again due to parallelism, this is deemed a successful result.

* Add timeout to auth rotation test
---
 internal/daemon/worker/auth_rotation_test.go  | 23 +++++++++++++++++--
 .../tests/api/credentials/credentials_test.go |  4 ++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/internal/daemon/worker/auth_rotation_test.go b/internal/daemon/worker/auth_rotation_test.go
index 6e4d1ac13e..3f4deee90a 100644
--- a/internal/daemon/worker/auth_rotation_test.go
+++ b/internal/daemon/worker/auth_rotation_test.go
@@ -4,6 +4,8 @@
 package worker_test
 
 import (
+	"bytes"
+	"context"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -109,12 +111,23 @@ func TestRotationTicking(t *testing.T) {
 
 	// Now we wait and check that we see new values in the DB and different
 	// creds on the worker after each rotation period
-	for i := 2; i < 5; i++ {
+
+	// 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)
+	defer deadlineCtxCancel()
+
+	rotationCount := 2
+	for {
+		if rotationCount > 5 {
+			break
+		}
 		nextRotation := w.Worker().AuthRotationNextRotation.Load()
 		time.Sleep((*nextRotation).Sub(time.Now()) + 5*time.Second)
 
 		// Verify we see 2- after credentials have rotated, we should see current and previous
-		auths, err = workerAuthRepo.List(c.Context(), (*types.NodeInformation)(nil))
+		auths, err = workerAuthRepo.List(deadlineCtx, (*types.NodeInformation)(nil))
 		require.NoError(err)
 		assert.Len(auths, 2)
 
@@ -126,6 +139,11 @@ func TestRotationTicking(t *testing.T) {
 			nodeenrollment.WithStorageWrapper(w.Config().WorkerAuthStorageKms),
 		)
 		require.NoError(err)
+		if bytes.Equal(currKey, currNodeCreds.CertificatePublicKeyPkix) {
+			// Load due to parallel tests may mean that we didn't hit rotation
+			// yet, loop back to evaluate again
+			continue
+		}
 		assert.NotEqual(currKey, currNodeCreds.CertificatePublicKeyPkix)
 		currKey = currNodeCreds.CertificatePublicKeyPkix
 		currKeyId, err := nodeenrollment.KeyIdFromPkix(currNodeCreds.CertificatePublicKeyPkix)
@@ -155,5 +173,6 @@ func TestRotationTicking(t *testing.T) {
 		require.NotNil(w.Worker().GrpcClientConn)
 		require.NoError(w.Worker().GrpcClientConn.Close())
 		require.NoError(w.Worker().StartControllerConnections())
+		rotationCount++
 	}
 }
diff --git a/internal/tests/api/credentials/credentials_test.go b/internal/tests/api/credentials/credentials_test.go
index 5dd32040e3..397eeed2ee 100644
--- a/internal/tests/api/credentials/credentials_test.go
+++ b/internal/tests/api/credentials/credentials_test.go
@@ -468,7 +468,7 @@ func TestUpdateAfterKeyRotation(t *testing.T) {
 
 	// Update JSON credential, will re-encrypt with new key versions
 	obj["password"] = "password"
-	cred, err = credsClient.Update(ctx, cred.Item.Id, 1, credentials.WithJsonCredentialObject(obj))
+	_, err = credsClient.Update(ctx, cred.Item.Id, 1, credentials.WithJsonCredentialObject(obj))
 	require.NoError(err)
 
 	// Create new key versions again
@@ -500,7 +500,7 @@ func TestUpdateAfterKeyRotation(t *testing.T) {
 	for {
 		jobs, err := scopesClient.ListKeyVersionDestructionJobs(ctx, proj.PublicId)
 		require.NoError(err)
-		if len(jobs.Items) == 1 {
+		if len(jobs.Items) >= 1 {
 			break
 		}
 		select {