From 24ee9995594930e2c6b180611f06f6b891b78cbe Mon Sep 17 00:00:00 2001 From: kxue43 Date: Thu, 18 Dec 2025 12:27:43 -0500 Subject: [PATCH 1/2] test: Add case of cache hits for `auth`. --- auth/command_test.go | 228 ++++++++++++++++++++++++++++--------------- 1 file changed, 150 insertions(+), 78 deletions(-) diff --git a/auth/command_test.go b/auth/command_test.go index 6606a41..335646e 100644 --- a/auth/command_test.go +++ b/auth/command_test.go @@ -64,13 +64,13 @@ func (m *HomeDirMocker) TearDown(t *testing.T) { } func TestAssumeRoleCmdRun(t *testing.T) { + aesKey, _, err := generateKey(keySize) + require.NoError(t, err, "should be able to generate a random encryption key") + fromKeyring = func() ([]byte, error) { t.Helper() - key, _, err := generateKey(keySize) - require.NoError(t, err, "should be able to generate a random encryption key") - - return key, nil + return aesKey, nil } defer func() { fromKeyring = keyringGet }() @@ -80,110 +80,182 @@ func TestAssumeRoleCmdRun(t *testing.T) { hdm.SetUp(t) defer hdm.TearDown(t) - // Create an AssumeRoleCmd with some fields already filled with valid values. - // We don't test CLI parsing in unit tests. - cmd := AssumeRoleCmd{ - MFASerial: "mfa-serial", - Profile: "profile", - Region: "us-east-1", - RoleSessionName: "ToolkitCLI", - DurationSeconds: 3600, - } - roleArn := "role-arn" - token := "123456" + expiration := time.Now().Add(10 * time.Hour) - expiration := time.Now() + var duration int32 = 3600 - duration := int32(cmd.DurationSeconds) + stubber := testtools.NewStubber() + defer testtools.ExitTest(stubber, t) + + t.Run("Happy path no cache", func(t *testing.T) { + // Create an AssumeRoleCmd with some fields already filled with valid values. + // We don't test CLI parsing in unit tests. + cmd := AssumeRoleCmd{ + MFASerial: "mfa-serial", + Profile: "profile", + Region: "us-east-1", + RoleSessionName: "ToolkitCLI", + DurationSeconds: int64(duration), + } + + token := "123456" + + mockedTtyDevice := &MockFileDescriptor{} + + _, err := mockedTtyDevice.r.WriteString(token + "\n") + require.NoError(t, err, "should be able to write token to mocked TTY file descriptor") + + tty := NewTTY(mockedTtyDevice, "toolkit-assume-role: ", 0) + + dest := MockFileDescriptor{} + + soutput := CredentialProcessOutput{ + AccessKeyId: "access-key-id", + SecretAccessKey: "secret-access-key", + SessionToken: "session-token", + Expiration: expiration.Format(time.RFC3339), + Version: 1, + } + + stubber.Add(testtools.Stub{ + OperationName: "AssumeRole", + Input: &sts.AssumeRoleInput{ + DurationSeconds: &duration, + RoleArn: &roleArn, + RoleSessionName: &cmd.RoleSessionName, + SerialNumber: &cmd.MFASerial, + TokenCode: &token, + }, + Output: &sts.AssumeRoleOutput{ + Credentials: &types.Credentials{ + AccessKeyId: &soutput.AccessKeyId, + SecretAccessKey: &soutput.SecretAccessKey, + SessionToken: &soutput.SessionToken, + Expiration: &expiration, + }, + }, + Error: nil, + }) - mockedTtyDevice := &MockFileDescriptor{} + ctx := context.Background() - _, err := mockedTtyDevice.r.WriteString(token + "\n") - require.NoError(t, err, "should be able to write token to mocked TTY file descriptor") + err = cmd.ValidateInputs([]string{roleArn}) + require.NoError(t, err, "fields of AssumeRoleCmd should validate without error") - tty := NewTTY(mockedTtyDevice, "toolkit-assume-role: ", 0) + // Use a mocked credentials provider so that we can load config without error during tests. + mockCredsProvider := credentials.NewStaticCredentialsProvider("dummy", "dummy", "dummy") - dest := MockFileDescriptor{} + cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(mockCredsProvider), config.WithRegion(cmd.Region)) + require.NoError(t, err, "should be able to load config using a mocked credentials provider") - soutput := CredentialProcessOutput{ - AccessKeyId: "access-key-id", - SecretAccessKey: "secret-access-key", - SessionToken: "session-token", - Expiration: expiration.Format(time.RFC3339), - Version: 1, - } + err = cmd.Init(tty, cfg) + require.NoError(t, err, "should be able to init caching without error") - stubber := testtools.NewStubber() + // Stub the STS client object. + cmd.client = sts.NewFromConfig(*stubber.SdkConfig) - stubber.Add(testtools.Stub{ - OperationName: "AssumeRole", - Input: &sts.AssumeRoleInput{ - DurationSeconds: &duration, - RoleArn: &roleArn, - RoleSessionName: &cmd.RoleSessionName, - SerialNumber: &cmd.MFASerial, - TokenCode: &token, - }, - Output: &sts.AssumeRoleOutput{ - Credentials: &types.Credentials{ - AccessKeyId: &soutput.AccessKeyId, - SecretAccessKey: &soutput.SecretAccessKey, - SessionToken: &soutput.SessionToken, - Expiration: &expiration, - }, - }, - Error: nil, + err = cmd.Run(ctx, &dest) + require.NoError(t, err, "should be able to run command without error") + + cacheFilePath := filepath.Join(hdm.TempDir, ".aws", "toolkit-cache", EncodeToFileName(roleArn, expiration)) + + info, err := os.Stat(cacheFilePath) + require.NoError(t, err, "should be able to locate the cache file created by the Run method") + + assert.False(t, info.IsDir(), "the cache file created by the Run method should be a regular file, not a directory") + + rawContents, err := os.ReadFile(filepath.Clean(cacheFilePath)) + require.NoError(t, err, "should be able to read cache file raw content without error") + + rawContents, err = cmd.cacher.cipher.Decrypt(rawContents) + require.NoError(t, err, "should be able to decrypt cache file without error") + + var sCachedContents CredentialProcessOutput + + err = json.Unmarshal(rawContents, &sCachedContents) + require.NoError(t, err, "should be able to unmarshal decrypted cache file without error") + + assert.Equal(t, sCachedContents.AccessKeyId, soutput.AccessKeyId, "AccessKeyId from cache file should match STS call result") + + assert.Equal(t, sCachedContents.SecretAccessKey, soutput.SecretAccessKey, "SecretAccessKey from cache file should match STS call result") + + assert.Equal(t, sCachedContents.SessionToken, soutput.SessionToken, "SessionToken from cache file should match STS call result") + + assert.Equal(t, sCachedContents.Expiration, soutput.Expiration, "Expiration from cache file should match STS call result") + + assert.Equal(t, sCachedContents.Version, soutput.Version, "Version from cache file should be the right value of 1") + + assert.Equal(t, rawContents, dest.w.Bytes(), "outputs to stdout should be identical to decrypted cache file contents") }) - ctx := context.Background() + t.Run("Happy path cache hits", func(t *testing.T) { + cmd := AssumeRoleCmd{ + MFASerial: "mfa-serial", + Profile: "profile", + Region: "us-east-1", + RoleSessionName: "ToolkitCLI", + DurationSeconds: int64(duration), + } + + mockedTtyDevice := &MockFileDescriptor{} + + tty := NewTTY(mockedTtyDevice, "toolkit-assume-role: ", 0) + + dest := MockFileDescriptor{} - err = cmd.ValidateInputs([]string{roleArn}) - require.NoError(t, err, "fields of AssumeRoleCmd should validate without error") + ctx := context.Background() - // Use a mocked credentials provider so that we can load config without error during tests. - mockCredsProvider := credentials.NewStaticCredentialsProvider("dummy", "dummy", "dummy") + err := cmd.ValidateInputs([]string{roleArn}) + require.NoError(t, err, "fields of AssumeRoleCmd should validate without error") - cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(mockCredsProvider), config.WithRegion(cmd.Region)) - require.NoError(t, err, "should be able to load config using a mocked credentials provider") + // Use a mocked credentials provider so that we can load config without error during tests. + mockCredsProvider := credentials.NewStaticCredentialsProvider("dummy", "dummy", "dummy") - err = cmd.Init(tty, cfg) - require.NoError(t, err, "should be able to init caching without error") + cfg, err := config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(mockCredsProvider), config.WithRegion(cmd.Region)) + require.NoError(t, err, "should be able to load config using a mocked credentials provider") - // Stub the STS client object. - cmd.client = sts.NewFromConfig(*stubber.SdkConfig) + err = cmd.Init(tty, cfg) + require.NoError(t, err, "should be able to init caching without error") - err = cmd.Run(ctx, &dest) - require.NoError(t, err, "should be able to run command without error") + // Stub the STS client object. When cache hits, the STS client shouldn't have been used at all. + cmd.client = nil - cacheFilePath := filepath.Join(hdm.TempDir, ".aws", "toolkit-cache", EncodeToFileName(roleArn, expiration)) + err = cmd.Run(ctx, &dest) + require.NoError(t, err, "should be able to run command without error") - info, err := os.Stat(cacheFilePath) - require.NoError(t, err, "should be able to locate the cache file created by the Run method") + cacheFilePath := filepath.Join(hdm.TempDir, ".aws", "toolkit-cache", EncodeToFileName(roleArn, expiration)) - assert.False(t, info.IsDir(), "the cache file created by the Run method should be a regular file, not a directory") + info, err := os.Stat(cacheFilePath) + require.NoError(t, err, "should be able to locate the cache file created by the Run method") - rawContents, err := os.ReadFile(filepath.Clean(cacheFilePath)) - require.NoError(t, err, "should be able to read cache file raw content without error") + assert.False(t, info.IsDir(), "the cache file created by the Run method should be a regular file, not a directory") - rawContents, err = cmd.cacher.cipher.Decrypt(rawContents) - require.NoError(t, err, "should be able to decrypt cache file without error") + rawContents, err := os.ReadFile(filepath.Clean(cacheFilePath)) + require.NoError(t, err, "should be able to read cache file raw content without error") - var sCachedContents CredentialProcessOutput + rawContents, err = cmd.cacher.cipher.Decrypt(rawContents) + require.NoError(t, err, "should be able to decrypt cache file without error") - err = json.Unmarshal(rawContents, &sCachedContents) - require.NoError(t, err, "should be able to unmarshal decrypted cache file without error") + var sCachedContents CredentialProcessOutput - assert.Equal(t, sCachedContents.AccessKeyId, soutput.AccessKeyId, "AccessKeyId from cache file should match STS call result") + err = json.Unmarshal(rawContents, &sCachedContents) + require.NoError(t, err, "should be able to unmarshal decrypted cache file without error") - assert.Equal(t, sCachedContents.SecretAccessKey, soutput.SecretAccessKey, "SecretAccessKey from cache file should match STS call result") + var stdoutContents CredentialProcessOutput - assert.Equal(t, sCachedContents.SessionToken, soutput.SessionToken, "SessionToken from cache file should match STS call result") + err = json.Unmarshal(dest.w.Bytes(), &stdoutContents) + require.NoError(t, err, "should be able to unmarshal outputs to stdout without error") - assert.Equal(t, sCachedContents.Expiration, soutput.Expiration, "Expiration from cache file should match STS call result") + assert.Equal(t, sCachedContents.AccessKeyId, stdoutContents.AccessKeyId, "AccessKeyId from cache file should match that from stdout") - assert.Equal(t, sCachedContents.Version, soutput.Version, "Version from cache file should be the right value of 1") + assert.Equal(t, sCachedContents.SecretAccessKey, stdoutContents.SecretAccessKey, "SecretAccessKey from cache file should match that from stdout") - assert.Equal(t, rawContents, dest.w.Bytes(), "outputs to stdout should be identical to decrypted cache file contents") + assert.Equal(t, sCachedContents.SessionToken, stdoutContents.SessionToken, "SessionToken from cache file should match that from stdout") + + assert.Equal(t, sCachedContents.Expiration, stdoutContents.Expiration, "Expiration from cache file should match that from stdout") + + assert.Equal(t, sCachedContents.Version, stdoutContents.Version, "Version from cache file should match that from stdout") + }) } From 9e596f29fbe81f62d029ba65ad05df16046f91ef Mon Sep 17 00:00:00 2001 From: kxue43 Date: Thu, 18 Dec 2025 12:48:17 -0500 Subject: [PATCH 2/2] refactor: Minor. --- auth/command_test.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/auth/command_test.go b/auth/command_test.go index 335646e..9973eee 100644 --- a/auth/command_test.go +++ b/auth/command_test.go @@ -86,9 +86,6 @@ func TestAssumeRoleCmdRun(t *testing.T) { var duration int32 = 3600 - stubber := testtools.NewStubber() - defer testtools.ExitTest(stubber, t) - t.Run("Happy path no cache", func(t *testing.T) { // Create an AssumeRoleCmd with some fields already filled with valid values. // We don't test CLI parsing in unit tests. @@ -119,6 +116,9 @@ func TestAssumeRoleCmdRun(t *testing.T) { Version: 1, } + stubber := testtools.NewStubber() + defer testtools.ExitTest(stubber, t) + stubber.Add(testtools.Stub{ OperationName: "AssumeRole", Input: &sts.AssumeRoleInput{