chore: introduce special logic for unit tests that rely on feature fl…

…ags (#12645)

CDKv1 introduced feature flags that allowed users to opt-in to new
backwards-incompatible behaviour. There are two types of flags -
one set that need to be entirely removed in CDKv2 and their
'enabled' behaviour made the default, and another set where the
the new behaviour should be made the default, and allow users to
explicitly opt-out.

This change addresses the testing strategy for the former set (and
not the latter). There exist unit tests that test the behaviour both
when the feature flags are enabled and feature flags and disabled.
However, in v2, the feature flags will be removed and its usage
blocked. The default behaviour will be as if the feature flag is
enabled.

Introduce branching logic that treats these unit tests depending on
whether it's executed in the context of CDKv1 or CDKv2.

The logic is as follows:
- In the context of v1, all tests will execute as normal.
- In the context of v2, tests that rely on feature flag to be unset will
   not be executed.
- In the context of v2, tests that rely on feature flag to be set will
  not configure the feature flag on the App's context and continue to
  execute the test.

This logic has been captured at a single location in two methods -
`testFeatureFlag()` and `testFeatureFlagDisabled()`.

To validate this logic, the unit tests that rely on the feature flags
`aws-kms:defaultKeyPolicies`,
`aws-secretsmanager:parseOwnedSecretName` and
`core:enableStackNameDuplicates` have been updated to use
these methods.

Forward looking: The final PR in v2 is expected to look like
this - #12644

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

packages/@aws-cdk/aws-kms/test/key.test.ts

@@ -2,6 +2,7 @@ import { arrayWith, ResourcePart } from '@aws-cdk/assert';
 import '@aws-cdk/assert/jest';
 import * as iam from '@aws-cdk/aws-iam';
 import * as cdk from '@aws-cdk/core';
+import { testFutureBehavior, testLegacyBehavior } from 'cdk-build-tools/lib/feature-flag';
 import * as kms from '../lib';
 
 const ADMIN_ACTIONS: string[] = [
@@ -39,19 +40,10 @@ const LEGACY_ADMIN_ACTIONS: string[] = [
   'kms:UntagResource',
 ];
 
-let app: cdk.App;
-let stack: cdk.Stack;
-beforeEach(() => {
-  app = new cdk.App({
-    context: {
-      // By default, enable the correct key policy behavior. Specific tests will test the disabled behavior.
-      '@aws-cdk/aws-kms:defaultKeyPolicies': true,
-    },
-  });
-  stack = new cdk.Stack(app);
-});
+const flags = { '@aws-cdk/aws-kms:defaultKeyPolicies': true };
 
-test('default key', () => {
+testFutureBehavior('default key', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   new kms.Key(stack, 'MyKey');
 
   expect(stack).toHaveResource('AWS::KMS::Key', {
@@ -75,14 +67,16 @@ test('default key', () => {
   }, ResourcePart.CompleteDefinition);
 });
 
-test('default with no retention', () => {
+testFutureBehavior('default with no retention', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   new kms.Key(stack, 'MyKey', { removalPolicy: cdk.RemovalPolicy.DESTROY });
 
   expect(stack).toHaveResource('AWS::KMS::Key', { DeletionPolicy: 'Delete', UpdateReplacePolicy: 'Delete' }, ResourcePart.CompleteDefinition);
 });
 
 describe('key policies', () => {
-  test('can specify a default key policy', () => {
+  testFutureBehavior('can specify a default key policy', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const policy = new iam.PolicyDocument();
     const statement = new iam.PolicyStatement({ resources: ['*'], actions: ['kms:Put*'] });
     statement.addArnPrincipal('arn:aws:iam::111122223333:root');
@@ -107,7 +101,8 @@ describe('key policies', () => {
     });
   });
 
-  test('can append to the default key policy', () => {
+  testFutureBehavior('can append to the default key policy', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const statement = new iam.PolicyStatement({ resources: ['*'], actions: ['kms:Put*'] });
     statement.addArnPrincipal('arn:aws:iam::111122223333:root');
 
@@ -139,16 +134,53 @@ describe('key policies', () => {
     });
   });
 
-  test.each([
-    ['decrypt', (key: kms.Key, user: iam.IGrantable) => key.grantDecrypt(user), 'kms:Decrypt'],
-    ['encrypt', (key: kms.Key, user: iam.IGrantable) => key.grantEncrypt(user), ['kms:Encrypt', 'kms:ReEncrypt*', 'kms:GenerateDataKey*']],
-  ])('grant %s', (_, grantFn, actions) => {
+  testFutureBehavior('decrypt', flags, cdk.App, (app) => {
+    // GIVEN
+    const stack = new cdk.Stack(app);
+    const key = new kms.Key(stack, 'Key');
+    const user = new iam.User(stack, 'User');
+
+    // WHEN
+    key.grantDecrypt(user);
+
+    // THEN
+    // Key policy should be unmodified by the grant.
+    expect(stack).toHaveResource('AWS::KMS::Key', {
+      KeyPolicy: {
+        Statement: [
+          {
+            Action: 'kms:*',
+            Effect: 'Allow',
+            Principal: { AWS: { 'Fn::Join': ['', ['arn:', { Ref: 'AWS::Partition' }, ':iam::', { Ref: 'AWS::AccountId' }, ':root']] } },
+            Resource: '*',
+          },
+        ],
+        Version: '2012-10-17',
+      },
+    });
+
+    expect(stack).toHaveResource('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [
+          {
+            Action: 'kms:Decrypt',
+            Effect: 'Allow',
+            Resource: { 'Fn::GetAtt': ['Key961B73FD', 'Arn'] },
+          },
+        ],
+        Version: '2012-10-17',
+      },
+    });
+  });
+
+  testFutureBehavior('encrypt', flags, cdk.App, (app) => {
     // GIVEN
+    const stack = new cdk.Stack(app);
     const key = new kms.Key(stack, 'Key');
     const user = new iam.User(stack, 'User');
 
     // WHEN
-    grantFn(key, user);
+    key.grantEncrypt(user);
 
     // THEN
     // Key policy should be unmodified by the grant.
@@ -170,7 +202,7 @@ describe('key policies', () => {
       PolicyDocument: {
         Statement: [
           {
-            Action: actions,
+            Action: ['kms:Encrypt', 'kms:ReEncrypt*', 'kms:GenerateDataKey*'],
             Effect: 'Allow',
             Resource: { 'Fn::GetAtt': ['Key961B73FD', 'Arn'] },
           },
@@ -180,7 +212,7 @@ describe('key policies', () => {
     });
   });
 
-  test('grant for a principal in a dependent stack works correctly', () => {
+  testFutureBehavior('grant for a principal in a dependent stack works correctly', flags, cdk.App, (app) => {
     const principalStack = new cdk.Stack(app, 'PrincipalStack');
     const principal = new iam.Role(principalStack, 'Role', {
       assumedBy: new iam.AnyPrincipal(),
@@ -213,7 +245,8 @@ describe('key policies', () => {
     });
   });
 
-  test('additional key admins can be specified (with imported/immutable principal)', () => {
+  testFutureBehavior('additional key admins can be specified (with imported/immutable principal)', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const adminRole = iam.Role.fromRoleArn(stack, 'Admin', 'arn:aws:iam::123456789012:role/TrustedAdmin');
     new kms.Key(stack, 'MyKey', { admins: [adminRole] });
 
@@ -242,7 +275,8 @@ describe('key policies', () => {
     });
   });
 
-  test('additional key admins can be specified (with owned/mutable principal)', () => {
+  testFutureBehavior('additional key admins can be specified (with owned/mutable principal)', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const adminRole = new iam.Role(stack, 'AdminRole', {
       assumedBy: new iam.AccountRootPrincipal(),
     });
@@ -279,7 +313,8 @@ describe('key policies', () => {
   });
 });
 
-test('key with some options', () => {
+testFutureBehavior('key with some options', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey', {
     enableKeyRotation: true,
     enabled: false,
@@ -311,17 +346,20 @@ test('key with some options', () => {
   });
 });
 
-test('setting pendingWindow value to not in allowed range will throw', () => {
+testFutureBehavior('setting pendingWindow value to not in allowed range will throw', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   expect(() => new kms.Key(stack, 'MyKey', { enableKeyRotation: true, pendingWindow: cdk.Duration.days(6) }))
     .toThrow('\'pendingWindow\' value must between 7 and 30 days. Received: 6');
 });
 
-test('setting trustAccountIdentities to false will throw (when the defaultKeyPolicies feature flag is enabled)', () => {
+testFutureBehavior('setting trustAccountIdentities to false will throw (when the defaultKeyPolicies feature flag is enabled)', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   expect(() => new kms.Key(stack, 'MyKey', { trustAccountIdentities: false }))
     .toThrow('`trustAccountIdentities` cannot be false if the @aws-cdk/aws-kms:defaultKeyPolicies feature flag is set');
 });
 
-test('addAlias creates an alias', () => {
+testFutureBehavior('addAlias creates an alias', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey', {
     enableKeyRotation: true,
     enabled: false,
@@ -342,7 +380,8 @@ test('addAlias creates an alias', () => {
   });
 });
 
-test('can run multiple addAlias', () => {
+testFutureBehavior('can run multiple addAlias', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey', {
     enableKeyRotation: true,
     enabled: false,
@@ -374,7 +413,8 @@ test('can run multiple addAlias', () => {
   });
 });
 
-test('keyId resolves to a Ref', () => {
+testFutureBehavior('keyId resolves to a Ref', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey');
 
   new cdk.CfnOutput(stack, 'Out', {
@@ -387,7 +427,8 @@ test('keyId resolves to a Ref', () => {
   });
 });
 
-test('fails if key policy has no actions', () => {
+testFutureBehavior('fails if key policy has no actions', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey');
 
   key.addToResourcePolicy(new iam.PolicyStatement({
@@ -398,7 +439,8 @@ test('fails if key policy has no actions', () => {
   expect(() => app.synth()).toThrow(/A PolicyStatement must specify at least one \'action\' or \'notAction\'/);
 });
 
-test('fails if key policy has no IAM principals', () => {
+testFutureBehavior('fails if key policy has no IAM principals', flags, cdk.App, (app) => {
+  const stack = new cdk.Stack(app);
   const key = new kms.Key(stack, 'MyKey');
 
   key.addToResourcePolicy(new iam.PolicyStatement({
@@ -410,14 +452,15 @@ test('fails if key policy has no IAM principals', () => {
 });
 
 describe('imported keys', () => {
-  test('throw an error when providing something that is not a valid key ARN', () => {
+  testFutureBehavior('throw an error when providing something that is not a valid key ARN', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     expect(() => {
       kms.Key.fromKeyArn(stack, 'Imported', 'arn:aws:kms:us-east-1:123456789012:key');
     }).toThrow(/KMS key ARN must be in the format 'arn:aws:kms:<region>:<account>:key\/<keyId>', got: 'arn:aws:kms:us-east-1:123456789012:key'/);
 
   });
 
-  test('can have aliases added to them', () => {
+  testFutureBehavior('can have aliases added to them', flags, cdk.App, (app) => {
     const stack2 = new cdk.Stack(app, 'Stack2');
     const myKeyImported = kms.Key.fromKeyArn(stack2, 'MyKeyImported',
       'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012');
@@ -443,15 +486,17 @@ describe('imported keys', () => {
 
 describe('addToResourcePolicy allowNoOp and there is no policy', () => {
   // eslint-disable-next-line jest/expect-expect
-  test('succeed if set to true (default)', () => {
+  testFutureBehavior('succeed if set to true (default)', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const key = kms.Key.fromKeyArn(stack, 'Imported',
       'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012');
 
     key.addToResourcePolicy(new iam.PolicyStatement({ resources: ['*'], actions: ['*'] }));
 
   });
 
-  test('fails if set to false', () => {
+  testFutureBehavior('fails if set to false', flags, cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const key = kms.Key.fromKeyArn(stack, 'Imported',
       'arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012');
 
@@ -463,16 +508,8 @@ describe('addToResourcePolicy allowNoOp and there is no policy', () => {
 });
 
 describe('when the defaultKeyPolicies feature flag is disabled', () => {
-  beforeEach(() => {
-    app = new cdk.App({
-      context: {
-        '@aws-cdk/aws-kms:defaultKeyPolicies': false,
-      },
-    });
-    stack = new cdk.Stack(app);
-  });
-
-  test('default key policy', () => {
+  testLegacyBehavior('default key policy', cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     new kms.Key(stack, 'MyKey');
 
     expect(stack).toHaveResource('AWS::KMS::Key', {
@@ -496,7 +533,8 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
     }, ResourcePart.CompleteDefinition);
   });
 
-  test('policy if specified appends to the default key policy', () => {
+  testLegacyBehavior('policy if specified appends to the default key policy', cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const key = new kms.Key(stack, 'MyKey');
     const p = new iam.PolicyStatement({ resources: ['*'], actions: ['kms:Encrypt'] });
     p.addArnPrincipal('arn:aws:iam::111122223333:root');
@@ -536,7 +574,8 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
     });
   });
 
-  test('trustAccountIdentities changes key policy to allow IAM control', () => {
+  testLegacyBehavior('trustAccountIdentities changes key policy to allow IAM control', cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     new kms.Key(stack, 'MyKey', { trustAccountIdentities: true });
     expect(stack).toHaveResourceLike('AWS::KMS::Key', {
       KeyPolicy: {
@@ -554,7 +593,8 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
     });
   });
 
-  test('additional key admins can be specified (with imported/immutable principal)', () => {
+  testLegacyBehavior('additional key admins can be specified (with imported/immutable principal)', cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const adminRole = iam.Role.fromRoleArn(stack, 'Admin', 'arn:aws:iam::123456789012:role/TrustedAdmin');
     new kms.Key(stack, 'MyKey', { admins: [adminRole] });
 
@@ -583,7 +623,8 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
     });
   });
 
-  test('additional key admins can be specified (with owned/mutable principal)', () => {
+  testLegacyBehavior('additional key admins can be specified (with owned/mutable principal)', cdk.App, (app) => {
+    const stack = new cdk.Stack(app);
     const adminRole = new iam.Role(stack, 'AdminRole', {
       assumedBy: new iam.AccountRootPrincipal(),
     });
@@ -627,8 +668,9 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
   });
 
   describe('grants', () => {
-    test('grant decrypt on a key', () => {
+    testLegacyBehavior('grant decrypt on a key', cdk.App, (app) => {
       // GIVEN
+      const stack = new cdk.Stack(app);
       const key = new kms.Key(stack, 'Key');
       const user = new iam.User(stack, 'User');
 
@@ -672,7 +714,7 @@ describe('when the defaultKeyPolicies feature flag is disabled', () => {
       });
     });
 
-    test('grant for a principal in a dependent stack works correctly', () => {
+    testLegacyBehavior('grant for a principal in a dependent stack works correctly', cdk.App, (app) => {
       const principalStack = new cdk.Stack(app, 'PrincipalStack');
       const principal = new iam.Role(principalStack, 'Role', {
         assumedBy: new iam.AnyPrincipal(),