From f89179d7db588b2f8156cef76aa0edf912845f93 Mon Sep 17 00:00:00 2001
From: James Kwon <96548424+hongil0316@users.noreply.github.com>
Date: Thu, 5 Dec 2024 22:44:17 -0500
Subject: [PATCH] do not log ec2 keypair

---
 modules/test-structure/save_test_data.go      | 15 ++++++-
 modules/test-structure/save_test_data_test.go | 42 +++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/modules/test-structure/save_test_data.go b/modules/test-structure/save_test_data.go
index 8f54975d3..50a2c67a9 100644
--- a/modules/test-structure/save_test_data.go
+++ b/modules/test-structure/save_test_data.go
@@ -65,7 +65,7 @@ func formatPackerOptionsPath(testFolder string) string {
 // SaveEc2KeyPair serializes and saves an Ec2KeyPair into the given folder. This allows you to create an Ec2KeyPair during setup
 // and to reuse that Ec2KeyPair later during validation and teardown.
 func SaveEc2KeyPair(t testing.TestingT, testFolder string, keyPair *aws.Ec2Keypair) {
-	SaveTestData(t, formatEc2KeyPairPath(testFolder), true, keyPair)
+	saveTestData(t, formatEc2KeyPairPath(testFolder), true, keyPair, false)
 }
 
 // LoadEc2KeyPair loads and unserializes an Ec2KeyPair from the given folder. This allows you to reuse an Ec2KeyPair that was
@@ -193,6 +193,15 @@ func FormatTestDataPath(testFolder string, filename string) string {
 // any contents that exist in the file found at `path` will be overwritten. This has the potential for causing duplicated resources
 // and should be used with caution. If `overwrite` is `false`, the save will be skipped and a warning will be logged.
 func SaveTestData(t testing.TestingT, path string, overwrite bool, value interface{}) {
+	saveTestData(t, path, overwrite, value, true)
+}
+
+// saveTestData serializes and saves a value used at test time to the given path. This allows you to create some sort of test data
+// (e.g., TerraformOptions) during setup and to reuse this data later during validation and teardown. If `overwrite` is `true`,
+// any contents that exist in the file found at `path` will be overwritten. This has the potential for causing duplicated resources
+// and should be used with caution. If `overwrite` is `false`, the save will be skipped and a warning will be logged.
+// If `loggedVal` is `true`, the value will be logged as JSON.
+func saveTestData(t testing.TestingT, path string, overwrite bool, value interface{}, loggedVal bool) {
 	logger.Default.Logf(t, "Storing test data in %s so it can be reused later", path)
 
 	if IsTestDataPresent(t, path) {
@@ -209,7 +218,9 @@ func SaveTestData(t testing.TestingT, path string, overwrite bool, value interfa
 		t.Fatalf("Failed to convert value %s to JSON: %v", path, err)
 	}
 
-	logger.Default.Logf(t, "Marshalled JSON: %s", string(bytes))
+	if loggedVal {
+		logger.Default.Logf(t, "Marshalled JSON: %s", string(bytes))
+	}
 
 	parentDir := filepath.Dir(path)
 	if err := os.MkdirAll(parentDir, 0777); err != nil {
diff --git a/modules/test-structure/save_test_data_test.go b/modules/test-structure/save_test_data_test.go
index 5e93ca0b6..9237088ca 100644
--- a/modules/test-structure/save_test_data_test.go
+++ b/modules/test-structure/save_test_data_test.go
@@ -1,13 +1,21 @@
 package test_structure
 
 import (
+	"encoding/json"
+	"fmt"
 	"os"
+	"strings"
 	"testing"
 
+	"github.com/gruntwork-io/terratest/modules/aws"
 	"github.com/gruntwork-io/terratest/modules/files"
 	"github.com/gruntwork-io/terratest/modules/k8s"
+	"github.com/gruntwork-io/terratest/modules/logger"
+	"github.com/gruntwork-io/terratest/modules/ssh"
 	"github.com/gruntwork-io/terratest/modules/terraform"
+	gotesting "github.com/gruntwork-io/terratest/modules/testing"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type testData struct {
@@ -272,3 +280,37 @@ func TestSaveAndLoadKubectlOptions(t *testing.T) {
 	actualData := LoadKubectlOptions(t, tmpFolder)
 	assert.Equal(t, expectedData, actualData)
 }
+
+type tStringLogger struct {
+	sb strings.Builder
+}
+
+func (l *tStringLogger) Logf(t gotesting.TestingT, format string, args ...interface{}) {
+	l.sb.WriteString(fmt.Sprintf(format, args...))
+	l.sb.WriteRune('\n')
+}
+
+func TestSaveAndLoadEC2KeyPair(t *testing.T) {
+	def, slogger := logger.Default, &tStringLogger{}
+	logger.Default = logger.New(slogger)
+	t.Cleanup(func() {
+		logger.Default = def
+	})
+
+	keyPair, err := ssh.GenerateRSAKeyPairE(t, 2048)
+	require.NoError(t, err)
+	ec2KeyPair := &aws.Ec2Keypair{
+		KeyPair: keyPair,
+		Name:    "test-ec2-key-pair",
+		Region:  "us-east-1",
+	}
+	storedEC2KeyPair, err := json.Marshal(ec2KeyPair)
+	require.NoError(t, err)
+
+	tmpFolder := t.TempDir()
+	SaveEc2KeyPair(t, tmpFolder, ec2KeyPair)
+	loadedEC2KeyPair := LoadEc2KeyPair(t, tmpFolder)
+	assert.Equal(t, ec2KeyPair, loadedEC2KeyPair)
+
+	assert.NotContains(t, slogger.sb.String(), string(storedEC2KeyPair), "stored ec2 key pair should not be logged")
+}