chore(cli-integ): optionally acquire environments from the cdk allocation service

packages/@aws-cdk-testing/cli-integ/lib/aws.ts

@@ -20,15 +20,19 @@ import { SNSClient } from '@aws-sdk/client-sns';
 import { SSOClient } from '@aws-sdk/client-sso';
 import { STSClient, GetCallerIdentityCommand } from '@aws-sdk/client-sts';
 import { fromIni } from '@aws-sdk/credential-providers';
-import type { AwsCredentialIdentityProvider } from '@smithy/types';
+import type { AwsCredentialIdentity, AwsCredentialIdentityProvider } from '@smithy/types';
 import { ConfiguredRetryStrategy } from '@smithy/util-retry';
 interface ClientConfig {
-  readonly credentials?: AwsCredentialIdentityProvider;
+  readonly credentials?: AwsCredentialIdentityProvider | AwsCredentialIdentity;
   readonly region: string;
   readonly retryStrategy: ConfiguredRetryStrategy;
 }
 
 export class AwsClients {
+  public static async forIdentity(region: string, identity: AwsCredentialIdentity, output: NodeJS.WritableStream) {
+    return new AwsClients(region, output, identity);
+  }
+
   public static async forRegion(region: string, output: NodeJS.WritableStream) {
     return new AwsClients(region, output);
   }
@@ -45,9 +49,12 @@ export class AwsClients {
   public readonly lambda: LambdaClient;
   public readonly sts: STSClient;
 
-  constructor(public readonly region: string, private readonly output: NodeJS.WritableStream) {
+  constructor(
+    public readonly region: string,
+    private readonly output: NodeJS.WritableStream,
+    public readonly identity?: AwsCredentialIdentity) {
     this.config = {
-      credentials: chainableCredentials(this.region),
+      credentials: this.identity ?? chainableCredentials(this.region),
       region: this.region,
       retryStrategy: new ConfiguredRetryStrategy(9, (attempt: number) => attempt ** 500),
     };

packages/@aws-cdk-testing/cli-integ/lib/integ-test.ts

@@ -13,6 +13,7 @@ if (SKIP_TESTS) {
 
 export interface TestContext {
   readonly randomString: string;
+  readonly name: string;
   readonly output: NodeJS.WritableStream;
   log(s: string): void;
 };
@@ -51,6 +52,7 @@ export function integTest(
       return await callback({
         output,
         randomString: randomString(),
+        name,
         log(s: string) {
           output.write(`${s}\n`);
         },

packages/@aws-cdk-testing/cli-integ/lib/with-aws.ts

@@ -1,8 +1,30 @@
+import { AtmosphereClient } from '@cdklabs/cdk-atmosphere-client';
 import { AwsClients } from './aws';
 import { TestContext } from './integ-test';
 import { ResourcePool } from './resource-pool';
 import { DisableBootstrapContext } from './with-cdk-app';
 
+export function atmosphereEnabled(): boolean {
+  const enabled = process.env.CDK_INTEG_ATMOSPHERE_ENABLED;
+  return enabled === 'true' || enabled === '1';
+}
+
+export function atmosphereEndpoint(): string {
+  const value = process.env.CDK_INTEG_ATMOSPHERE_ENDPOINT;
+  if (!value) {
+    throw new Error('CDK_INTEG_ATMOSPHERE_ENDPOINT is not defined');
+  }
+  return value;
+}
+
+export function atmospherePool() {
+  const value = process.env.CDK_INTEG_ATMOSPHERE_POOL;
+  if (!value) {
+    throw new Error('CDK_INTEG_ATMOSPHERE_POOL is not defined');
+  }
+  return value;
+}
+
 export type AwsContext = { readonly aws: AwsClients };
 
 /**
@@ -14,12 +36,40 @@ export function withAws<A extends TestContext>(
   block: (context: A & AwsContext & DisableBootstrapContext) => Promise<void>,
   disableBootstrap: boolean = false,
 ): (context: A) => Promise<void> {
-  return (context: A) => regionPool().using(async (region) => {
-    const aws = await AwsClients.forRegion(region, context.output);
-    await sanityCheck(aws);
+  return async (context: A) => {
+
+    if (atmosphereEnabled()) {
+      const atmosphere = new AtmosphereClient(atmosphereEndpoint());
+      const allocation = await atmosphere.acquire({ pool: atmospherePool(), requester: context.name });
+      const aws = await AwsClients.forIdentity(allocation.environment.region, {
+        accessKeyId: allocation.credentials.accessKeyId,
+        secretAccessKey: allocation.credentials.secretAccessKey,
+        sessionToken: allocation.credentials.sessionToken,
+        accountId: allocation.environment.account,
+      }, context.output);
+
+      await sanityCheck(aws);
+
+      let outcome = 'success';
+      try {
+        return await block({ ...context, disableBootstrap, aws });
+      } catch (e: any) {
+        outcome = 'failure';
+        throw e;
+      } finally {
+        await atmosphere.release(allocation.id, outcome);
+      }
+
+    } else {
+      return regionPool().using(async (region) => {
+        const aws = await AwsClients.forRegion(region, context.output);
+        await sanityCheck(aws);
+
+        return block({ ...context, disableBootstrap, aws });
+      });
+    }
 
-    return block({ ...context, disableBootstrap, aws });
-  });
+  };
 }
 
 let _regionPool: undefined | ResourcePool;

packages/@aws-cdk-testing/cli-integ/lib/with-cdk-app.ts

@@ -11,7 +11,7 @@ import { IPackageSource } from './package-sources/source';
 import { packageSourceInSubprocess } from './package-sources/subprocess';
 import { RESOURCES_DIR } from './resources';
 import { shell, ShellOptions, ShellHelper, rimraf } from './shell';
-import { AwsContext, withAws } from './with-aws';
+import { AwsContext, atmosphereEnabled, withAws } from './with-aws';
 import { withTimeout } from './with-timeout';
 
 export const DEFAULT_TEST_TIMEOUT_S = 20 * 60;
@@ -498,13 +498,22 @@ export class TestFixture extends ShellHelper {
 
     await this.packages.makeCliAvailable();
 
+    // if tests are using an explicit aws identity already (i.e creds)
+    // force every cdk command to use the same identity.
+    const awsCreds: Record<string, string> = this.aws.identity ? {
+      AWS_ACCESS_KEY_ID: this.aws.identity.accessKeyId,
+      AWS_SECRET_ACCESS_KEY: this.aws.identity.secretAccessKey,
+      AWS_SESSION_TOKEN: this.aws.identity.sessionToken!,
+    } : {};
+
     return this.shell(['cdk', ...(verbose ? ['-v'] : []), ...args], {
       ...options,
       modEnv: {
         AWS_REGION: this.aws.region,
         AWS_DEFAULT_REGION: this.aws.region,
         STACK_NAME_PREFIX: this.stackNamePrefix,
         PACKAGE_LAYOUT_VERSION: this.packages.majorVersion(),
+        ...awsCreds,
         ...options.modEnv,
       },
     });
@@ -555,37 +564,44 @@ export class TestFixture extends ShellHelper {
    * Cleanup leftover stacks and bootstrapped resources
    */
   public async dispose(success: boolean) {
-    const stacksToDelete = await this.deleteableStacks(this.stackNamePrefix);
-
-    this.sortBootstrapStacksToTheEnd(stacksToDelete);
-
-    // Bootstrap stacks have buckets that need to be cleaned
-    const bucketNames = stacksToDelete.map(stack => outputFromStack('BucketName', stack)).filter(defined);
-    // Parallelism will be reasonable
-    // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
-    await Promise.all(bucketNames.map(b => this.aws.emptyBucket(b)));
-    // The bootstrap bucket has a removal policy of RETAIN by default, so add it to the buckets to be cleaned up.
-    this.bucketsToDelete.push(...bucketNames);
-
-    // Bootstrap stacks have ECR repositories with images which should be deleted
-    const imageRepositoryNames = stacksToDelete.map(stack => outputFromStack('ImageRepositoryName', stack)).filter(defined);
-    // Parallelism will be reasonable
-    // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
-    await Promise.all(imageRepositoryNames.map(r => this.aws.deleteImageRepository(r)));
-
-    await this.aws.deleteStacks(
-      ...stacksToDelete.map((s) => {
-        if (!s.StackName) {
-          throw new Error('Stack name is required to delete a stack.');
-        }
-        return s.StackName;
-      }),
-    );
 
-    // We might have leaked some buckets by upgrading the bootstrap stack. Be
-    // sure to clean everything.
-    for (const bucket of this.bucketsToDelete) {
-      await this.aws.deleteBucket(bucket);
+    // when using the atmosphere service, it does resource cleanup on our behalf
+    // so we don't have to wait for it.
+    if (!atmosphereEnabled()) {
+
+      const stacksToDelete = await this.deleteableStacks(this.stackNamePrefix);
+
+      this.sortBootstrapStacksToTheEnd(stacksToDelete);
+
+      // Bootstrap stacks have buckets that need to be cleaned
+      const bucketNames = stacksToDelete.map(stack => outputFromStack('BucketName', stack)).filter(defined);
+      // Parallelism will be reasonable
+      // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
+      await Promise.all(bucketNames.map(b => this.aws.emptyBucket(b)));
+      // The bootstrap bucket has a removal policy of RETAIN by default, so add it to the buckets to be cleaned up.
+      this.bucketsToDelete.push(...bucketNames);
+
+      // Bootstrap stacks have ECR repositories with images which should be deleted
+      const imageRepositoryNames = stacksToDelete.map(stack => outputFromStack('ImageRepositoryName', stack)).filter(defined);
+      // Parallelism will be reasonable
+      // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
+      await Promise.all(imageRepositoryNames.map(r => this.aws.deleteImageRepository(r)));
+
+      await this.aws.deleteStacks(
+        ...stacksToDelete.map((s) => {
+          if (!s.StackName) {
+            throw new Error('Stack name is required to delete a stack.');
+          }
+          return s.StackName;
+        }),
+      );
+
+      // We might have leaked some buckets by upgrading the bootstrap stack. Be
+      // sure to clean everything.
+      for (const bucket of this.bucketsToDelete) {
+        await this.aws.deleteBucket(bucket);
+      }
+
     }
 
     // If the tests completed successfully, happily delete the fixture
@@ -662,7 +678,11 @@ export async function ensureBootstrapped(fixture: TestFixture) {
     },
   });
 
-  ALREADY_BOOTSTRAPPED_IN_THIS_RUN.add(envSpecifier);
+  // when using the atmosphere service, every test needs to bootstrap
+  // its own environment.
+  if (!atmosphereEnabled()) {
+    ALREADY_BOOTSTRAPPED_IN_THIS_RUN.add(envSpecifier);
+  }
 }
 
 function defined<A>(x: A): x is NonNullable<A> {

packages/@aws-cdk-testing/cli-integ/package.json

@@ -50,6 +50,7 @@
     "@aws-sdk/client-sso": "3.632.0",
     "@aws-sdk/client-sts": "3.632.0",
     "@aws-sdk/credential-providers": "3.632.0",
+    "@cdklabs/cdk-atmosphere-client": "0.0.1",
     "@smithy/util-retry": "3.0.8",
     "@smithy/types": "3.6.0",
     "axios": "1.7.7",
@@ -86,4 +87,4 @@
   "publishConfig": {
     "tag": "latest"
   }
-}
+}

yarn.lock

@@ -4612,7 +4612,7 @@
     "@smithy/types" "^3.7.1"
     tslib "^2.6.2"
 
-"@aws-sdk/credential-providers@^3.699.0":
+"@aws-sdk/credential-providers@^3.699.0", "@aws-sdk/credential-providers@^3.731.1":
   version "3.734.0"
   resolved "https://registry.npmjs.org/@aws-sdk/credential-providers/-/credential-providers-3.734.0.tgz#fae6d9ce10737c08da3f11f4864ca74ab98704b9"
   integrity sha512-3q76ngVxwX/kSRA0bjH7hUkIOVf/38aACmYpbwwr7jyRU3Cpbsj57W9YtRd7zS9/A4Jt6fYx7VFEA52ajyoGAQ==
@@ -6305,6 +6305,14 @@
   resolved "https://registry.npmjs.org/@bcoe/v8-coverage/-/v8-coverage-0.2.3.tgz#75a2e8b51cb758a7553d6804a5932d7aace75c39"
   integrity sha512-0hYQ8SB4Db5zvZB4axdMHGwEaQjkZzFjQiN9LVYvIFB2nSUHW9tYpxWriPrWDASIxiaXax83REcLxuSdnGPZtw==
 
+"@cdklabs/cdk-atmosphere-client@0.0.1":
+  version "0.0.1"
+  resolved "https://registry.npmjs.org/@cdklabs/cdk-atmosphere-client/-/cdk-atmosphere-client-0.0.1.tgz#48499daef9894a2905167ec27b5f34dd2ae523e0"
+  integrity sha512-y7kCj9ClOeMkd2iRTlp3hx+vTML4NtitI4MMogbN9iSsB8U+qdrXgu14LbKCAfQPRDoNbUYkgyZQwK4MkEP0Lw==
+  dependencies:
+    "@aws-sdk/credential-providers" "^3.731.1"
+    aws4fetch "^1.0.20"
+
 "@cdklabs/eslint-plugin@^1.3.0":
   version "1.3.2"
   resolved "https://registry.npmjs.org/@cdklabs/eslint-plugin/-/eslint-plugin-1.3.2.tgz#9a37485e0c94cd13a9becdd69791d4ff1dc1c515"
@@ -10144,6 +10152,11 @@ aws-sdk@^2.1379.0:
     uuid "8.0.0"
     xml2js "0.6.2"
 
+aws4fetch@^1.0.20:
+  version "1.0.20"
+  resolved "https://registry.npmjs.org/aws4fetch/-/aws4fetch-1.0.20.tgz#090d6c65e32c6df645dd5e5acf04cc56da575cbe"
+  integrity sha512-/djoAN709iY65ETD6LKCtyyEI04XIBP5xVvfmNxsEP0uJB5tyaGBztSryRr4HqMStr9R06PisQE7m9zDTXKu6g==
+
 axios@1.7.7:
   version "1.7.7"
   resolved "https://registry.npmjs.org/axios/-/axios-1.7.7.tgz#2f554296f9892a72ac8d8e4c5b79c14a91d0a47f"
@@ -19035,16 +19048,7 @@ string-length@^4.0.1:
     char-regex "^1.0.2"
     strip-ansi "^6.0.0"
 
-"string-width-cjs@npm:string-width@^4.2.0":
-  version "4.2.3"
-  resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010"
-  integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g==
-  dependencies:
-    emoji-regex "^8.0.0"
-    is-fullwidth-code-point "^3.0.0"
-    strip-ansi "^6.0.1"
-
-string-width@*, string-width@^1.0.1, "string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.3, string-width@^5.0.1, string-width@^5.1.2:
+"string-width-cjs@npm:string-width@^4.2.0", string-width@*, string-width@^1.0.1, "string-width@^1.0.2 || 2 || 3 || 4", string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.3, string-width@^5.0.1, string-width@^5.1.2:
   version "4.2.3"
   resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010"
   integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g==
@@ -19113,7 +19117,7 @@ stringify-package@^1.0.1:
   resolved "https://registry.npmjs.org/stringify-package/-/stringify-package-1.0.1.tgz#e5aa3643e7f74d0f28628b72f3dad5cecfc3ba85"
   integrity sha512-sa4DUQsYciMP1xhKWGuFM04fB0LG/9DlluZoSVywUMRNvzid6XucHK0/90xGxRoHrAaROrcHK1aPKaijCtSrhg==
 
-"strip-ansi-cjs@npm:strip-ansi@^6.0.1":
+"strip-ansi-cjs@npm:strip-ansi@^6.0.1", strip-ansi@^6.0.0, strip-ansi@^6.0.1:
   version "6.0.1"
   resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9"
   integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A==
@@ -19127,13 +19131,6 @@ strip-ansi@^3.0.1:
   dependencies:
     ansi-regex "^2.0.0"
 
-strip-ansi@^6.0.0, strip-ansi@^6.0.1:
-  version "6.0.1"
-  resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9"
-  integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A==
-  dependencies:
-    ansi-regex "^5.0.1"
-
 strip-ansi@^7.0.1, strip-ansi@^7.1.0:
   version "7.1.0"
   resolved "https://registry.npmjs.org/strip-ansi/-/strip-ansi-7.1.0.tgz#d5b6568ca689d8561370b0707685d22434faff45"
@@ -20214,7 +20211,7 @@ workerpool@^6.5.1:
   resolved "https://registry.npmjs.org/workerpool/-/workerpool-6.5.1.tgz#060f73b39d0caf97c6db64da004cd01b4c099544"
   integrity sha512-Fs4dNYcsdpYSAfVxhnl1L5zTksjvOJxtC5hzMNl+1t9B8hTJTdKDyZ5ju7ztgPy+ft9tBFXoOlDNiOT9WUXZlA==
 
-"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0":
+"wrap-ansi-cjs@npm:wrap-ansi@^7.0.0", wrap-ansi@^7.0.0:
   version "7.0.0"
   resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43"
   integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q==
@@ -20232,15 +20229,6 @@ wrap-ansi@^6.0.1, wrap-ansi@^6.2.0:
     string-width "^4.1.0"
     strip-ansi "^6.0.0"
 
-wrap-ansi@^7.0.0:
-  version "7.0.0"
-  resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-7.0.0.tgz#67e145cff510a6a6984bdf1152911d69d2eb9e43"
-  integrity sha512-YVGIj2kamLSTxw6NsZjoBxfSwsn0ycdesmc4p+Q21c5zPuZ1pl+NfxVdxPtdHvmNVOQ6XSYG4AUtyt/Fi7D16Q==
-  dependencies:
-    ansi-styles "^4.0.0"
-    string-width "^4.1.0"
-    strip-ansi "^6.0.0"
-
 wrap-ansi@^8.1.0:
   version "8.1.0"
   resolved "https://registry.npmjs.org/wrap-ansi/-/wrap-ansi-8.1.0.tgz#56dc22368ee570face1b49819975d9b9a5ead214"