Agent: Improve TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput (#26286)

* Improve TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput to make it more robust in race test

* Tweak the sensitivity on waiting for template re-renders after triggering

command/agent/agent_auto_auth_self_heal_test.go

@@ -18,9 +18,11 @@ import (
 	agentConfig "github.com/hashicorp/vault/command/agent/config"
 	"github.com/hashicorp/vault/command/agent/template"
 	"github.com/hashicorp/vault/command/agentproxyshared/auth"
-	token_file "github.com/hashicorp/vault/command/agentproxyshared/auth/token-file"
+	tokenfile "github.com/hashicorp/vault/command/agentproxyshared/auth/token-file"
 	"github.com/hashicorp/vault/command/agentproxyshared/sink"
 	"github.com/hashicorp/vault/command/agentproxyshared/sink/file"
+	"github.com/hashicorp/vault/helper/testhelpers/corehelpers"
+	"github.com/hashicorp/vault/helper/testhelpers/minimal"
 	vaulthttp "github.com/hashicorp/vault/http"
 	"github.com/hashicorp/vault/sdk/helper/logging"
 	"github.com/hashicorp/vault/sdk/helper/pointerutil"
@@ -46,58 +48,52 @@ path "/secret/*" {
 // if the token is revoked, Auto Auth is re-triggered and a valid new token
 // is written to a sink, and the template is correctly rendered with the new token
 func TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput(t *testing.T) {
-	logger := logging.NewVaultLogger(hclog.Trace)
+	// Unset the environment variable so that agent picks up the right test cluster address
-	cluster := vault.NewTestCluster(t,
+	t.Setenv(api.EnvVaultAddress, "")
-		&vault.CoreConfig{},
-		&vault.TestClusterOptions{
-			NumCores:    1,
-			HandlerFunc: vaulthttp.Handler,
-		})
-	cluster.Start()
-	defer cluster.Cleanup()
 
-	vault.TestWaitActive(t, cluster.Cores[0].Core)
+	tmpDir := t.TempDir()
+	pathLookupSelf := filepath.Join(tmpDir, "lookup-self")
+	pathVaultToken := filepath.Join(tmpDir, "vault-token")
+	pathTokenFile := filepath.Join(tmpDir, "token-file")
+
+	secretRenderInterval := 1 * time.Second
+	contextTimeout := 30 * time.Second
+
+	cluster := minimal.NewTestSoloCluster(t, nil)
+	logger := corehelpers.NewTestLogger(t)
 	serverClient := cluster.Cores[0].Client
 
-	// Unset the environment variable so that agent picks up the right test
-	// cluster address
-	defer os.Setenv(api.EnvVaultAddress, os.Getenv(api.EnvVaultAddress))
-	os.Unsetenv(api.EnvVaultAddress)
-
-	// create temp dir for this test run
-	tmpDir, err := os.MkdirTemp("", "TestAutoAuth_SelfHealing")
-	require.NoError(t, err)
-	defer os.RemoveAll(tmpDir)
-
 	// Create token
 	secret, err := serverClient.Auth().Token().Create(&api.TokenCreateRequest{
 		Policies: []string{"test-autoauth"},
 	})
 	require.NoError(t, err)
+	require.NotNil(t, secret)
+	require.NotNil(t, secret.Auth)
+	require.NotEmpty(t, secret.Auth.ClientToken)
 	token := secret.Auth.ClientToken
 
 	// Write token to vault-token file
-	tokenFilePath := filepath.Join(tmpDir, "vault-token")
+	tokenFile, err := os.Create(pathVaultToken)
-	tokenFile, err := os.Create(tokenFilePath)
 	require.NoError(t, err)
 	_, err = tokenFile.WriteString(token)
 	require.NoError(t, err)
 	err = tokenFile.Close()
 	require.NoError(t, err)
 
-	defer os.Remove(tokenFilePath)
+	// Give us some leeway of 3 errors 1 from each of: auth handler, sink server template server.
-	require.NoError(t, err)
+	errCh := make(chan error, 3)
-
+	ctx, cancel := context.WithTimeout(context.Background(), contextTimeout)
-	ctx, cancel := context.WithCancel(context.Background())
 
 	// Create auth handler
-	am, err := token_file.NewTokenFileAuthMethod(&auth.AuthConfig{
+	am, err := tokenfile.NewTokenFileAuthMethod(&auth.AuthConfig{
 		Logger: logger.Named("auth.method"),
 		Config: map[string]interface{}{
-			"token_file_path": filepath.Join(tmpDir, "vault-token"),
+			"token_file_path": pathVaultToken,
 		},
 	})
 	require.NoError(t, err)
+
 	ahConfig := &auth.AuthHandlerConfig{
 		Logger:                       logger.Named("auth.handler"),
 		Client:                       serverClient,
@@ -107,34 +103,20 @@ func TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput(t *testing.T) {
 		ExitOnError:                  false,
 	}
 	ah := auth.NewAuthHandler(ahConfig)
-	errCh := make(chan error)
-
 	go func() {
 		errCh <- ah.Run(ctx, am)
 	}()
-	defer func() {
-		select {
-		case <-ctx.Done():
-		case err := <-errCh:
-			if err != nil {
-				t.Fatal(err)
-			}
-		}
-	}()
 
 	// Create sink file server
-	sinkFilePath := filepath.Join(tmpDir, "token-file")
+	_, err = os.Create(pathTokenFile)
-	_, err = os.Create(sinkFilePath)
-	defer os.Remove(sinkFilePath)
 	require.NoError(t, err)
 
 	config := &sink.SinkConfig{
 		Logger: logger.Named("sink.file"),
 		Config: map[string]interface{}{
-			"path": sinkFilePath,
+			"path": pathTokenFile,
 		},
 	}
-
 	fs, err := file.NewFileSink(config)
 	if err != nil {
 		t.Fatal(err)
@@ -145,34 +127,24 @@ func TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput(t *testing.T) {
 		Logger: logger.Named("sink.server"),
 		Client: serverClient,
 	})
-
 	go func() {
 		errCh <- ss.Run(ctx, ah.OutputCh, []*sink.SinkConfig{config}, ah.AuthInProgress)
 	}()
-	defer func() {
-		select {
-		case <-ctx.Done():
-		case err := <-errCh:
-			if err != nil {
-				t.Fatal(err)
-			}
-		}
-	}()
 
 	// Create template server
-	sc := template.ServerConfig{
+	sc := &template.ServerConfig{
-		Logger: logging.NewVaultLogger(hclog.Trace),
+		Logger: logger.Named("template.server"),
 		AgentConfig: &agentConfig.Config{
 			Vault: &agentConfig.Vault{
 				Address:       serverClient.Address(),
 				TLSSkipVerify: true,
 			},
 			TemplateConfig: &agentConfig.TemplateConfig{
-				StaticSecretRenderInt: time.Second * 2,
+				StaticSecretRenderInt: secretRenderInterval,
 			},
 			AutoAuth: &agentConfig.AutoAuth{
 				Sinks: []*agentConfig.Sink{{Type: "file", Config: map[string]interface{}{
-					"path": filepath.Join(filepath.Join(tmpDir, "lookup-self")),
+					"path": pathLookupSelf,
 				}}},
 			},
 			ExitAfterAuth: false,
@@ -183,68 +155,77 @@ func TestAutoAuthSelfHealing_TokenFileAuth_SinkOutput(t *testing.T) {
 	}
 
 	templateTest := &ctconfig.TemplateConfig{
-		Contents: pointerutil.StringPtr(lookupSelfTemplateContents),
+		Contents:    pointerutil.StringPtr(lookupSelfTemplateContents),
+		Destination: pointerutil.StringPtr(pathLookupSelf),
 	}
-	dstFile := fmt.Sprintf("%s/%s", tmpDir, "lookup-self")
-	templateTest.Destination = pointerutil.StringPtr(dstFile)
 	templatesToRender := []*ctconfig.TemplateConfig{templateTest}
 
 	var server *template.Server
-	server = template.NewServer(&sc)
+	server = template.NewServer(sc)
-
 	go func() {
 		errCh <- server.Run(ctx, ah.TemplateTokenCh, templatesToRender, ah.AuthInProgress, ah.InvalidToken)
 	}()
-	defer func() {
-		select {
-		case <-ctx.Done():
-		case err := <-errCh:
-			if err != nil {
-				t.Fatal(err)
-			}
-		}
-	}()
 
-	// Must be done at the very end so that nothing is blocking
+	// Trigger template render (mark the time as being earlier, based on the render interval)
-	defer cancel()
+	preTriggerTime := time.Now().Add(-secretRenderInterval)
-
-	// Trigger template render
 	ah.TemplateTokenCh <- token
-	fileInfo, err := waitForFiles(t, filepath.Join(tmpDir, "token-file"), time.Time{})
+	fileInfo, err := waitForFiles(t, pathTokenFile, preTriggerTime)
 	require.NoError(t, err)
 
-	tokenInSink, err := os.ReadFile(filepath.Join(tmpDir, "token-file"))
+	tokenInSink, err := os.ReadFile(pathTokenFile)
 	require.NoError(t, err)
-	require.Equal(t, string(tokenInSink), token)
+	require.Equal(t, token, string(tokenInSink))
 
 	// Revoke Token
 	t.Logf("revoking token")
-	serverClient.Auth().Token().RevokeOrphan(token)
+	err = serverClient.Auth().Token().RevokeOrphan(token)
+	require.NoError(t, err)
 
 	// Create new token
-	tokenSecret, err := serverClient.Auth().Token().Create(&api.TokenCreateRequest{
+	tokenSecret, err := serverClient.Auth().Token().Create(&api.TokenCreateRequest{})
-		Policies: []string{"test-autoauth"},
-	})
 	require.NoError(t, err)
+	require.NotNil(t, tokenSecret)
+	require.NotNil(t, tokenSecret.Auth)
+	require.NotEmpty(t, tokenSecret.Auth.ClientToken)
 	newToken := tokenSecret.Auth.ClientToken
 
 	// Write token to file
-	err = os.WriteFile(filepath.Join(tmpDir, "vault-token"), []byte(newToken), 0o600)
+	err = os.WriteFile(pathVaultToken, []byte(newToken), 0o600)
 	require.NoError(t, err)
 
 	// Wait for auto-auth to complete
-	_, err = waitForFiles(t, filepath.Join(tmpDir, "token-file"), fileInfo.ModTime())
+	updatedFileInfo, err := waitForFiles(t, pathTokenFile, fileInfo.ModTime())
 	require.NoError(t, err)
 
 	// Verify the new token has been written to a file sink after re-authenticating using lookup-self
-	tokenInSink, err = os.ReadFile(filepath.Join(tmpDir, "token-file"))
+	tokenInSink, err = os.ReadFile(pathTokenFile)
+	require.NoError(t, err)
+	require.Equal(t, newToken, string(tokenInSink))
+
+	// Wait for the lookup-self file to be updated (again)
+	_, err = waitForFiles(t, pathLookupSelf, updatedFileInfo.ModTime())
 	require.NoError(t, err)
-	require.Equal(t, string(tokenInSink), newToken)
 
 	// Verify the template has now been correctly rendered with the new token
-	templateContents, err := os.ReadFile(filepath.Join(tmpDir, "lookup-self"))
+	templateContents, err := os.ReadFile(pathLookupSelf)
 	require.NoError(t, err)
-	require.Equal(t, string(templateContents), newToken)
+	require.Equal(t, newToken, string(templateContents))
+
+	// Calling cancel will stop the 'Run' funcs we started in Goroutines, we should
+	// then check that there were no errors in our channel.
+	cancel()
+	wrapUpTimeout := 5 * time.Second
+	for {
+		select {
+		case <-time.After(wrapUpTimeout):
+			t.Fatal("test timed out")
+		case err := <-errCh:
+			require.NoError(t, err)
+		case <-ctx.Done():
+			// We can finish the test ourselves
+			return
+		}
+	}
 }
 
 // Test_NoAutoAuthSelfHealing_BadPolicy tests that auto auth
@@ -297,7 +278,7 @@ func Test_NoAutoAuthSelfHealing_BadPolicy(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 
 	// Create auth handler
-	am, err := token_file.NewTokenFileAuthMethod(&auth.AuthConfig{
+	am, err := tokenfile.NewTokenFileAuthMethod(&auth.AuthConfig{
 		Logger: logger.Named("auth.method"),
 		Config: map[string]interface{}{
 			"token_file_path": filepath.Join(filepath.Join(tmpDir, "vault-token")),
@@ -451,6 +432,8 @@ func Test_NoAutoAuthSelfHealing_BadPolicy(t *testing.T) {
 }
 
 func waitForFiles(t *testing.T, filePath string, prevModTime time.Time) (os.FileInfo, error) {
+	t.Helper()
+
 	var err error
 	var fileInfo os.FileInfo
 	tick := time.Tick(100 * time.Millisecond)
@@ -459,7 +442,7 @@ func waitForFiles(t *testing.T, filePath string, prevModTime time.Time) (os.File
 	for {
 		select {
 		case <-timeout:
-			return nil, fmt.Errorf("timed out waiting for templates to render, last error: %v", err)
+			return nil, fmt.Errorf("timed out waiting for templates to render, last error: %w", err)
 		case <-tick:
 		}