From 474d869d2ec93bc8048bcf427ad9d61c3ed5ec39 Mon Sep 17 00:00:00 2001
From: William Will <10997562+willwill96@users.noreply.github.com>
Date: Fri, 8 Aug 2025 11:04:53 -0600
Subject: [PATCH] Add unit tests

---
 pkg/sessions/redis/aws-iam/auth.go      |  2 +-
 pkg/sessions/redis/aws-iam/auth_test.go | 35 +++++++++++++++++++++++
 pkg/sessions/redis/redis_store_test.go  | 38 +++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 pkg/sessions/redis/aws-iam/auth_test.go

diff --git a/pkg/sessions/redis/aws-iam/auth.go b/pkg/sessions/redis/aws-iam/auth.go
index b2387abf..8d86a1f2 100644
--- a/pkg/sessions/redis/aws-iam/auth.go
+++ b/pkg/sessions/redis/aws-iam/auth.go
@@ -42,7 +42,7 @@ type iamTokenGenerator struct {
 }
 
 // New creates a new IAMTokenGenerator instance
-func New(serviceName, clusterName, userName string) (*IAMTokenGenerator, error) {
+func New(serviceName, clusterName, userName string) (TokenGenerator, error) {
 
 	ctx := context.Background()
 	cfg, err := config.LoadDefaultConfig(ctx)
diff --git a/pkg/sessions/redis/aws-iam/auth_test.go b/pkg/sessions/redis/aws-iam/auth_test.go
new file mode 100644
index 00000000..3b19ff20
--- /dev/null
+++ b/pkg/sessions/redis/aws-iam/auth_test.go
@@ -0,0 +1,35 @@
+package auth
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestAWSIAMTokenGenerator(t *testing.T) {
+	// Set up the environment, so we don't make any external calls to AWS
+	t.Setenv("AWS_CONFIG_FILE", "file_not_exists")
+	t.Setenv("AWS_SHARED_CREDENTIALS_FILE", "file_not_exists")
+	t.Setenv("AWS_ENDPOINT_URL", "http://localhost:9999/aws")
+	t.Setenv("AWS_ACCESS_KEY_ID", "access_key")
+	t.Setenv("AWS_SECRET_ACCESS_KEY", "secret_key")
+	t.Setenv("AWS_REGION", "us-east-1")
+
+	tokenGenerator, err := New("elasticache", "test-cluster", "test-user")
+	require.NotNil(t, tokenGenerator)
+	require.NoError(t, err)
+
+	token, err := tokenGenerator.GenerateToken()
+	require.NoError(t, err)
+	require.NotEmpty(t, token)
+	require.Contains(t, token, "X-Amz-Algorithm", "signed token should contain algorithm attribute")
+	require.Contains(t, token, "User=test-user", "signed token should contain user parameter")
+	require.Contains(t, token, "X-Amz-Credential", "signed token should contain credential attribute")
+	require.Contains(t, token, "X-Amz-Date", "signed token should contain date attribute")
+	require.Contains(t, token, "X-Amz-Expires", "signed token should contain expires attribute")
+	require.Contains(t, token, "X-Amz-SignedHeaders", "signed token should contain signed headers attribute")
+	require.Contains(t, token, "X-Amz-Signature", "signed token should contain signature attribute")
+	require.Contains(t, token, "Action=connect", "signed token should contain connect action")
+	require.False(t, strings.HasPrefix(token, "http://"), "token should not have http:// scheme")
+}
diff --git a/pkg/sessions/redis/redis_store_test.go b/pkg/sessions/redis/redis_store_test.go
index 1bff6855..6edd5b2e 100644
--- a/pkg/sessions/redis/redis_store_test.go
+++ b/pkg/sessions/redis/redis_store_test.go
@@ -11,6 +11,7 @@ import (
 	"github.com/oauth2-proxy/oauth2-proxy/v7/pkg/sessions/tests"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/redis/go-redis/v9"
 )
 
 const (
@@ -271,4 +272,41 @@ var _ = Describe("Redis SessionStore Tests", func() {
 			Expect(opts).To(BeNil())
 		})
 	})
+
+	Describe("AWSIAMAuth", func() {
+		Context("with AWS IAM options", func() {
+			It("should initialize CredentialsProvider when AWSIAMConfig is present", func() {
+				redisOpts := options.RedisStoreOptions{
+					AWSIAMConfig: &options.AWSIAMOptions{
+						ServiceName: "elasticache",
+						ClusterName: "test-cluster",
+						Username:    "test-user",
+					},
+				}
+
+				var opt *redis.Options
+				opt = &redis.Options{}
+
+				err := setupAWSIAMAuth(redisOpts, opt)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(opt.CredentialsProvider).ToNot(BeNil())
+
+				// Verify the CredentialsProvider returns the expected username
+				username, _ := opt.CredentialsProvider()
+				Expect(username).To(Equal("test-user"))
+			})
+			It("should not initialize CredentialsProvider when AWSIAMConfig is nil", func() {
+				redisOpts := options.RedisStoreOptions{
+					AWSIAMConfig: nil,
+				}
+
+				var opt *redis.Options
+				opt = &redis.Options{}
+
+				err := setupAWSIAMAuth(redisOpts, opt)
+				Expect(err).ToNot(HaveOccurred())
+				Expect(opt.CredentialsProvider).To(BeNil())
+			})
+		})
+	})
 })