feat(gen2-migration): preserve existing stack policies in lock step (#14648)

* chore: update lock execute  and rollback

* chore: add unit tests

* chore: update tests

* chore: address lock idempotency issue

* chore: update lock tests

---------

Co-authored-by: Sai Ray <saisujit@amazon.com>

Author: sai-ray

packages/amplify-cli/src/__tests__/commands/gen2-migration/lock.test.ts

@@ -0,0 +1,257 @@
+import { AmplifyMigrationLockStep } from '../../../commands/gen2-migration/lock';
+import { $TSContext } from '@aws-amplify/amplify-cli-core';
+import { CloudFormationClient, SetStackPolicyCommand } from '@aws-sdk/client-cloudformation';
+import { AmplifyClient, UpdateAppCommand } from '@aws-sdk/client-amplify';
+import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
+import { SpinningLogger } from '../../../commands/gen2-migration/_spinning-logger';
+
+jest.mock('@aws-sdk/client-cloudformation', () => ({
+  ...jest.requireActual('@aws-sdk/client-cloudformation'),
+  CloudFormationClient: jest.fn(),
+}));
+jest.mock('@aws-sdk/client-amplify', () => ({
+  ...jest.requireActual('@aws-sdk/client-amplify'),
+  AmplifyClient: jest.fn(),
+}));
+jest.mock('@aws-sdk/client-appsync', () => ({
+  ...jest.requireActual('@aws-sdk/client-appsync'),
+  AppSyncClient: jest.fn().mockImplementation(() => ({ send: jest.fn() })),
+  paginateListGraphqlApis: jest.fn().mockImplementation(() => ({
+    [Symbol.asyncIterator]: async function* () {
+      yield { graphqlApis: [{ name: 'testApp-testEnv', apiId: 'test-api-id' }] };
+    },
+  })),
+}));
+jest.mock('@aws-sdk/client-dynamodb', () => ({
+  ...jest.requireActual('@aws-sdk/client-dynamodb'),
+  DynamoDBClient: jest.fn().mockImplementation(() => ({ send: jest.fn() })),
+  paginateListTables: jest.fn().mockImplementation(() => ({
+    [Symbol.asyncIterator]: async function* () {
+      yield { TableNames: ['Table1-test-api-id-testEnv', 'Table2-test-api-id-testEnv'] };
+    },
+  })),
+}));
+jest.mock('@aws-amplify/amplify-prompts', () => ({
+  printer: { info: jest.fn(), blankLine: jest.fn(), success: jest.fn(), warn: jest.fn(), debug: jest.fn() },
+  AmplifySpinner: jest.fn().mockImplementation(() => ({
+    start: jest.fn(),
+    stop: jest.fn(),
+    resetMessage: jest.fn(),
+  })),
+  isDebug: false,
+}));
+jest.mock('../../../commands/gen2-migration/_validations', () => ({
+  AmplifyGen2MigrationValidations: jest.fn().mockImplementation(() => ({
+    validateDeploymentStatus: jest.fn().mockResolvedValue(undefined),
+    validateDrift: jest.fn().mockResolvedValue(undefined),
+  })),
+}));
+
+describe('AmplifyMigrationLockStep', () => {
+  let lockStep: AmplifyMigrationLockStep;
+  let mockCfnSend: jest.Mock;
+  let mockAmplifySend: jest.Mock;
+  let mockLogger: SpinningLogger;
+
+  beforeEach(() => {
+    mockCfnSend = jest.fn();
+    mockAmplifySend = jest.fn();
+
+    (CloudFormationClient as jest.Mock).mockImplementation(() => ({ send: mockCfnSend }));
+    (AmplifyClient as jest.Mock).mockImplementation(() => ({ send: mockAmplifySend }));
+    (DynamoDBClient as jest.Mock).mockImplementation(() => ({ send: jest.fn() }));
+
+    mockLogger = new SpinningLogger('mock');
+    jest.spyOn(mockLogger, 'info').mockImplementation(() => {});
+    jest.spyOn(mockLogger, 'start').mockImplementation(() => {});
+    jest.spyOn(mockLogger, 'succeed').mockImplementation(() => {});
+    jest.spyOn(mockLogger, 'push').mockImplementation(() => {});
+    jest.spyOn(mockLogger, 'pop').mockImplementation(() => {});
+
+    lockStep = new AmplifyMigrationLockStep(
+      mockLogger,
+      'testEnv',
+      'testApp',
+      'test-app-id',
+      'test-root-stack',
+      'us-east-1',
+      {} as $TSContext,
+    );
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  describe('forward stack policy merge', () => {
+    it('should append lock statement to empty stack policy', async () => {
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: undefined }).mockResolvedValueOnce({});
+      mockAmplifySend.mockResolvedValueOnce({ app: { environmentVariables: {} } }).mockResolvedValueOnce({});
+
+      const plan = await lockStep.forward();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(1);
+      expect(setCalls[0][0].input).toEqual({
+        StackName: 'test-root-stack',
+        StackPolicyBody: JSON.stringify({
+          Statement: [{ Effect: 'Deny', Action: 'Update:*', Principal: '*', Resource: '*' }],
+        }),
+      });
+    });
+
+    it('should append lock statement preserving existing statements', async () => {
+      const existingPolicy = {
+        Statement: [{ Effect: 'Deny', Action: 'Update:Replace', Principal: '*', Resource: 'LogicalResourceId/MyDB' }],
+      };
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: JSON.stringify(existingPolicy) }).mockResolvedValueOnce({});
+      mockAmplifySend.mockResolvedValueOnce({ app: { environmentVariables: {} } }).mockResolvedValueOnce({});
+
+      const plan = await lockStep.forward();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(1);
+      expect(setCalls[0][0].input).toEqual({
+        StackName: 'test-root-stack',
+        StackPolicyBody: JSON.stringify({
+          Statement: [
+            { Effect: 'Deny', Action: 'Update:Replace', Principal: '*', Resource: 'LogicalResourceId/MyDB' },
+            { Effect: 'Deny', Action: 'Update:*', Principal: '*', Resource: '*' },
+          ],
+        }),
+      });
+    });
+
+    it('should skip SetStackPolicy when lock statement already exists', async () => {
+      const alreadyLockedPolicy = {
+        Statement: [{ Effect: 'Deny', Action: 'Update:*', Principal: '*', Resource: '*' }],
+      };
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: JSON.stringify(alreadyLockedPolicy) });
+      mockAmplifySend.mockResolvedValueOnce({ app: { environmentVariables: {} } }).mockResolvedValueOnce({});
+
+      const plan = await lockStep.forward();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(0);
+    });
+  });
+
+  describe('forward env var merge', () => {
+    it('should merge new env var with existing env vars', async () => {
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: undefined }).mockResolvedValueOnce({});
+      mockAmplifySend.mockResolvedValueOnce({ app: { environmentVariables: { EXISTING: 'value' } } }).mockResolvedValueOnce({});
+
+      const plan = await lockStep.forward();
+      await plan.execute();
+
+      const updateCalls = mockAmplifySend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof UpdateAppCommand);
+      expect(updateCalls).toHaveLength(1);
+      expect(updateCalls[0][0].input).toEqual({
+        appId: 'test-app-id',
+        environmentVariables: { EXISTING: 'value', GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv' },
+      });
+    });
+  });
+
+  describe('rollback stack policy removal', () => {
+    it('should remove lock statement and preserve customer statements', async () => {
+      const policyWithLock = {
+        Statement: [
+          { Effect: 'Deny', Action: 'Update:Replace', Principal: '*', Resource: 'LogicalResourceId/MyDB' },
+          { Effect: 'Deny', Action: 'Update:*', Principal: '*', Resource: '*' },
+        ],
+      };
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: JSON.stringify(policyWithLock) }).mockResolvedValueOnce({});
+      mockAmplifySend
+        .mockResolvedValueOnce({ app: { environmentVariables: { GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv' } } })
+        .mockResolvedValueOnce({});
+
+      const plan = await lockStep.rollback();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(1);
+      expect(setCalls[0][0].input).toEqual({
+        StackName: 'test-root-stack',
+        StackPolicyBody: JSON.stringify({
+          Statement: [{ Effect: 'Deny', Action: 'Update:Replace', Principal: '*', Resource: 'LogicalResourceId/MyDB' }],
+        }),
+      });
+    });
+
+    it('should set allow-all when lock statement was the only one', async () => {
+      const policyWithOnlyLock = {
+        Statement: [{ Effect: 'Deny', Action: 'Update:*', Principal: '*', Resource: '*' }],
+      };
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: JSON.stringify(policyWithOnlyLock) }).mockResolvedValueOnce({});
+      mockAmplifySend
+        .mockResolvedValueOnce({ app: { environmentVariables: { GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv' } } })
+        .mockResolvedValueOnce({});
+
+      const plan = await lockStep.rollback();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(1);
+      expect(setCalls[0][0].input).toEqual({
+        StackName: 'test-root-stack',
+        StackPolicyBody: JSON.stringify({
+          Statement: [{ Effect: 'Allow', Action: 'Update:*', Principal: '*', Resource: '*' }],
+        }),
+      });
+    });
+
+    it('should skip SetStackPolicy when no existing policy (lock not found)', async () => {
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: undefined });
+      mockAmplifySend
+        .mockResolvedValueOnce({ app: { environmentVariables: { GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv' } } })
+        .mockResolvedValueOnce({});
+
+      const plan = await lockStep.rollback();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(0);
+    });
+
+    it('should skip SetStackPolicy when lock statement is not found', async () => {
+      const customerPolicy = {
+        Statement: [{ Effect: 'Deny', Action: 'Update:Replace', Principal: '*', Resource: 'LogicalResourceId/MyDB' }],
+      };
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: JSON.stringify(customerPolicy) });
+      mockAmplifySend
+        .mockResolvedValueOnce({ app: { environmentVariables: { GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv' } } })
+        .mockResolvedValueOnce({});
+
+      const plan = await lockStep.rollback();
+      await plan.execute();
+
+      const setCalls = mockCfnSend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof SetStackPolicyCommand);
+      expect(setCalls).toHaveLength(0);
+    });
+  });
+
+  describe('rollback env var removal', () => {
+    it('should remove GEN2_MIGRATION_ENVIRONMENT_NAME and preserve other env vars', async () => {
+      mockCfnSend.mockResolvedValueOnce({ StackPolicyBody: undefined });
+      mockAmplifySend
+        .mockResolvedValueOnce({
+          app: { environmentVariables: { GEN2_MIGRATION_ENVIRONMENT_NAME: 'testEnv', OTHER: 'keep' } },
+        })
+        .mockResolvedValueOnce({});
+
+      const plan = await lockStep.rollback();
+      await plan.execute();
+
+      const updateCalls = mockAmplifySend.mock.calls.filter(([cmd]: [unknown]) => cmd instanceof UpdateAppCommand);
+      expect(updateCalls).toHaveLength(1);
+      expect(updateCalls[0][0].input).toEqual({
+        appId: 'test-app-id',
+        environmentVariables: { OTHER: 'keep' },
+      });
+    });
+  });
+});