From bb30e79ffe8a05ba0e1401fa53257fbaff5706a3 Mon Sep 17 00:00:00 2001
From: Otavio Macedo <288203+otaviomacedo@users.noreply.github.com>
Date: Thu, 14 Nov 2024 13:14:48 +0000
Subject: [PATCH] More unit tests

---
 .../lib/api/aws-auth/awscli-compatible.ts     |  2 +-
 .../api/aws-auth/awscli-compatible.test.ts    | 58 ++++++++++++++-----
 2 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/packages/aws-cdk/lib/api/aws-auth/awscli-compatible.ts b/packages/aws-cdk/lib/api/aws-auth/awscli-compatible.ts
index 6f2500c5e48ba..146aa4ac3c57a 100644
--- a/packages/aws-cdk/lib/api/aws-auth/awscli-compatible.ts
+++ b/packages/aws-cdk/lib/api/aws-auth/awscli-compatible.ts
@@ -182,7 +182,7 @@ function caBundlePathFromEnvironment(): string | undefined {
 function shouldPrioritizeEnv() {
   const id = process.env.AWS_ACCESS_KEY_ID || process.env.AMAZON_ACCESS_KEY_ID;
   const key = process.env.AWS_SECRET_ACCESS_KEY || process.env.AMAZON_SECRET_ACCESS_KEY;
- const sessionToken = process.env.AWS_SESSION_TOKEN || process.env.AMAZON_SESSION_TOKEN;
+  const sessionToken = process.env.AWS_SESSION_TOKEN || process.env.AMAZON_SESSION_TOKEN;
   if (sessionToken) {
     process.env.AWS_SESSION_TOKEN = sessionToken;
   }
diff --git a/packages/aws-cdk/test/api/aws-auth/awscli-compatible.test.ts b/packages/aws-cdk/test/api/aws-auth/awscli-compatible.test.ts
index 64470214dfcbd..7e11b83f468df 100644
--- a/packages/aws-cdk/test/api/aws-auth/awscli-compatible.test.ts
+++ b/packages/aws-cdk/test/api/aws-auth/awscli-compatible.test.ts
@@ -1,26 +1,52 @@
 import { AwsCliCompatible } from '../../../lib/api/aws-auth/awscli-compatible';
 
-test('does not mess up with session token env variables if they are undefined', async () => {
-  process.env.AWS_ACCESS_KEY_ID = 'foo';
-  process.env.SECRET_ACCESS_KEY = 'bar';
+describe('Session token', () => {
+  test('does not mess up with session token env variables if they are undefined', async () => {
+    process.env.AWS_ACCESS_KEY_ID = 'foo';
+    process.env.SECRET_ACCESS_KEY = 'bar';
 
-  // Making sure these variables are not defined
-  delete process.env.AWS_SESSION_TOKEN;
-  delete process.env.AMAZON_SESSION_TOKEN;
+    // Making sure these variables are not defined
+    delete process.env.AWS_SESSION_TOKEN;
+    delete process.env.AMAZON_SESSION_TOKEN;
 
-  await AwsCliCompatible.credentialChainBuilder();
+    await AwsCliCompatible.credentialChainBuilder();
 
-  expect(process.env.AWS_SESSION_TOKEN).toBeUndefined();
-});
+    expect(process.env.AWS_SESSION_TOKEN).toBeUndefined();
+  });
 
-test('preserves session token env variables if they are defined', async () => {
-  process.env.AWS_ACCESS_KEY_ID = 'foo';
-  process.env.SECRET_ACCESS_KEY = 'bar';
+  test('preserves AWS_SESSION_TOKEN if it is defined', async () => {
+    process.env.AWS_ACCESS_KEY_ID = 'foo';
+    process.env.SECRET_ACCESS_KEY = 'bar';
 
-  process.env.AWS_SESSION_TOKEN = 'aaa';
-  process.env.AMAZON_SESSION_TOKEN = 'bbb';
+    process.env.AWS_SESSION_TOKEN = 'aaa';
+    delete process.env.AMAZON_SESSION_TOKEN;
 
-  await AwsCliCompatible.credentialChainBuilder();
+    await AwsCliCompatible.credentialChainBuilder();
 
-  expect(process.env.AWS_SESSION_TOKEN).toEqual('aaa');
+    expect(process.env.AWS_SESSION_TOKEN).toEqual('aaa');
+  });
+
+  test('assigns AWS_SESSION_TOKEN if it is not defined but AMAZON_SESSION_TOKEN is', async () => {
+    process.env.AWS_ACCESS_KEY_ID = 'foo';
+    process.env.SECRET_ACCESS_KEY = 'bar';
+
+    delete process.env.AWS_SESSION_TOKEN;
+    process.env.AMAZON_SESSION_TOKEN = 'aaa';
+
+    await AwsCliCompatible.credentialChainBuilder();
+
+    expect(process.env.AWS_SESSION_TOKEN).toEqual('aaa');
+  });
+
+  test('preserves AWS_SESSION_TOKEN if both are defined', async () => {
+    process.env.AWS_ACCESS_KEY_ID = 'foo';
+    process.env.SECRET_ACCESS_KEY = 'bar';
+
+    process.env.AWS_SESSION_TOKEN = 'aaa';
+    process.env.AMAZON_SESSION_TOKEN = 'bbb';
+
+    await AwsCliCompatible.credentialChainBuilder();
+
+    expect(process.env.AWS_SESSION_TOKEN).toEqual('aaa');
+  });
 });
\ No newline at end of file