refactor(ami-housekeeper): delegate SSM batch fetching to shared getParameters utility

Author: thomasnemer

lambdas/functions/ami-housekeeper/src/ami.test.ts

@@ -10,16 +10,18 @@ import {
 import {
   DescribeParametersCommand,
   DescribeParametersCommandOutput,
-  GetParametersCommand,
   SSMClient,
 } from '@aws-sdk/client-ssm';
+import { getParameters } from '@aws-github-runner/aws-ssm-util';
 import { mockClient } from 'aws-sdk-client-mock';
 import 'aws-sdk-client-mock-jest/vitest';
 
 import { AmiCleanupOptions, amiCleanup, defaultAmiCleanupOptions } from './ami';
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { fail } from 'assert';
 
+vi.mock('@aws-github-runner/aws-ssm-util');
+
 process.env.AWS_REGION = 'eu-east-1';
 const deleteAmisOlderThenDays = 30;
 const date31DaysAgo = new Date(new Date().setDate(new Date().getDate() - (deleteAmisOlderThenDays + 1)));
@@ -83,23 +85,12 @@ describe("delete AMI's", () => {
     mockSSMClient.reset();
 
     mockSSMClient.on(DescribeParametersCommand).resolves(ssmParameters);
-    mockSSMClient.on(GetParametersCommand).resolves({
-      Parameters: [
-        {
-          Name: 'ami-id/ami-ssm0001',
-          Type: 'String',
-          Value: 'ami-ssm0001',
-          Version: 1,
-        },
-        {
-          Name: 'ami-id/ami-ssm0002',
-          Type: 'String',
-          Value: 'ami-ssm0002',
-          Version: 1,
-        },
-      ],
-      InvalidParameters: [],
-    });
+    vi.mocked(getParameters).mockResolvedValue(
+      new Map([
+        ['ami-id/ami-ssm0001', 'ami-ssm0001'],
+        ['ami-id/ami-ssm0002', 'ami-ssm0002'],
+      ]),
+    );
 
     mockEC2Client.on(DescribeLaunchTemplatesCommand).resolves({
       LaunchTemplates: [
@@ -144,11 +135,7 @@ describe("delete AMI's", () => {
     expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplatesCommand);
     expect(mockEC2Client).toHaveReceivedCommand(DescribeLaunchTemplateVersionsCommand);
     expect(mockSSMClient).toHaveReceivedCommand(DescribeParametersCommand);
-    expect(mockSSMClient).toHaveReceivedCommandTimes(GetParametersCommand, 1);
-    expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-      Names: ['ami-id/ami-ssm0001', 'ami-id/ami-ssm0002'],
-      WithDecryption: true,
-    });
+    expect(getParameters).toHaveBeenCalledWith(['ami-id/ami-ssm0001', 'ami-id/ami-ssm0002']);
   });
 
   it('should NOT delete instances in use.', async () => {
@@ -484,17 +471,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand).resolves({
-        Parameters: [
-          {
-            Name: '/github-runner/config/ami_id',
-            Type: 'String',
-            Value: 'ami-underscore0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
+      vi.mocked(getParameters).mockResolvedValue(
+        new Map([['/github-runner/config/ami_id', 'ami-underscore0001']]),
+      );
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -503,10 +482,7 @@ describe("delete AMI's", () => {
 
       // AMI should not be deleted because it's referenced in SSM
       expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/github-runner/config/ami_id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
     });
 
@@ -521,17 +497,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand).resolves({
-        Parameters: [
-          {
-            Name: '/github-runner/config/ami-id',
-            Type: 'String',
-            Value: 'ami-hyphen0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
+      vi.mocked(getParameters).mockResolvedValue(
+        new Map([['/github-runner/config/ami-id', 'ami-hyphen0001']]),
+      );
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -540,10 +508,7 @@ describe("delete AMI's", () => {
 
       // AMI should not be deleted because it's referenced in SSM
       expect(mockEC2Client).not.toHaveReceivedCommand(DeregisterImageCommand);
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/github-runner/config/ami-id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami-id']);
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
     });
 
@@ -568,17 +533,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand).resolves({
-        Parameters: [
-          {
-            Name: '/some/path/ami-id',
-            Type: 'String',
-            Value: 'ami-wildcard0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
+      vi.mocked(getParameters).mockResolvedValue(
+        new Map([['/some/path/ami-id', 'ami-wildcard0001']]),
+      );
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -590,10 +547,7 @@ describe("delete AMI's", () => {
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/some/path/ami-id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/some/path/ami-id']);
     });
 
     it('handles wildcard SSM parameter patterns (*ami_id)', async () => {
@@ -617,17 +571,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand).resolves({
-        Parameters: [
-          {
-            Name: '/github-runner/config/ami_id',
-            Type: 'String',
-            Value: 'ami-wildcard-underscore0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
+      vi.mocked(getParameters).mockResolvedValue(
+        new Map([['/github-runner/config/ami_id', 'ami-wildcard-underscore0001']]),
+      );
 
       await amiCleanup({
         minimumDaysOld: 0,
@@ -639,10 +585,7 @@ describe("delete AMI's", () => {
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami_id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/github-runner/config/ami_id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/github-runner/config/ami_id']);
     });
 
     it('handles mixed explicit names and wildcard patterns', async () => {
@@ -664,17 +607,9 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand, { Names: ['/explicit/ami_id'], WithDecryption: true }).resolves({
-        Parameters: [
-          {
-            Name: '/explicit/ami_id',
-            Type: 'String',
-            Value: 'ami-explicit0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
+      vi.mocked(getParameters)
+        .mockResolvedValueOnce(new Map([['/explicit/ami_id', 'ami-explicit0001']]))
+        .mockResolvedValueOnce(new Map([['/discovered/ami-id', 'ami-wildcard0001']]));
 
       mockSSMClient.on(DescribeParametersCommand).resolves({
         Parameters: [
@@ -686,18 +621,6 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand, { Names: ['/discovered/ami-id'], WithDecryption: true }).resolves({
-        Parameters: [
-          {
-            Name: '/discovered/ami-id',
-            Type: 'String',
-            Value: 'ami-wildcard0001',
-            Version: 1,
-          },
-        ],
-        InvalidParameters: [],
-      });
-
       await amiCleanup({
         minimumDaysOld: 0,
         ssmParameterNames: ['/explicit/ami_id', '*ami-id'],
@@ -709,17 +632,11 @@ describe("delete AMI's", () => {
         ImageId: 'ami-unused0001',
       });
 
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/explicit/ami_id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/explicit/ami_id']);
       expect(mockSSMClient).toHaveReceivedCommandWith(DescribeParametersCommand, {
         ParameterFilters: [{ Key: 'Name', Option: 'Contains', Values: ['ami-id'] }],
       });
-      expect(mockSSMClient).toHaveReceivedCommandWith(GetParametersCommand, {
-        Names: ['/discovered/ami-id'],
-        WithDecryption: true,
-      });
+      expect(getParameters).toHaveBeenCalledWith(['/discovered/ami-id']);
     });
 
     it('handles SSM parameter fetch failures gracefully', async () => {
@@ -733,7 +650,7 @@ describe("delete AMI's", () => {
         ],
       });
 
-      mockSSMClient.on(GetParametersCommand).rejects(new Error('ParameterNotFound'));
+      vi.mocked(getParameters).mockRejectedValue(new Error('ParameterNotFound'));
 
       // Should not throw and should delete the AMI since SSM reference failed
       await amiCleanup({
@@ -791,7 +708,7 @@ describe("delete AMI's", () => {
         ImageId: 'ami-no-ssm0001',
       });
       expect(mockSSMClient).not.toHaveReceivedCommand(DescribeParametersCommand);
-      expect(mockSSMClient).not.toHaveReceivedCommand(GetParametersCommand);
+      expect(getParameters).not.toHaveBeenCalled();
     });
   });
 });

lambdas/functions/ami-housekeeper/src/ami.ts

@@ -8,9 +8,10 @@ import {
   Filter,
   Image,
 } from '@aws-sdk/client-ec2';
-import { GetParametersCommand, SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
+import { SSMClient, DescribeParametersCommand } from '@aws-sdk/client-ssm';
 import { createChildLogger } from '@aws-github-runner/aws-powertools-util';
 import { getTracedAWSV3Client } from '@aws-github-runner/aws-powertools-util';
+import { getParameters } from '@aws-github-runner/aws-ssm-util';
 
 const logger = createChildLogger('ami');
 
@@ -184,56 +185,37 @@ async function deleteSnapshot(options: AmiCleanupOptions, amiDetails: Image, ec2
 }
 
 /**
- * Resolves the values of multiple SSM parameters by their names in batched API calls.
+ * Resolves the values of multiple SSM parameters by their names.
+ * Delegates batching to the shared `getParameters` utility.
  * Doesn't fail on errors, but warns instead, as this process is best-effort.
  *
  * @param names - Array of SSM parameter names to resolve
- * @param ssmClient - Configured SSM client for making API calls
- * @returns Array of parameter values (may contain undefined for missing/failed parameters)
+ * @returns Array of parameter values in the same order as input (undefined for missing/failed parameters)
  */
-async function resolveSsmParameterValues(names: string[], ssmClient: SSMClient): Promise<(string | undefined)[]> {
+async function resolveSsmParameterValues(names: string[]): Promise<(string | undefined)[]> {
   if (names.length === 0) {
     return [];
   }
 
-  const results = new Map<string, string | undefined>();
-
-  // AWS SSM GetParameters API has a limit of 10 parameters per call
-  const chunkSize = 10;
-  for (let i = 0; i < names.length; i += chunkSize) {
-    const chunk = names.slice(i, i + chunkSize);
-    try {
-      const response = await ssmClient.send(
-        new GetParametersCommand({
-          Names: chunk,
-          WithDecryption: true,
-        }),
-      );
-
-      for (const param of response.Parameters ?? []) {
-        if (param.Name) {
-          results.set(param.Name, param.Value);
-        }
-      }
+  try {
+    const parameterMap = await getParameters(names);
 
-      // Log warnings for invalid parameters
-      for (const invalidParam of response.InvalidParameters ?? []) {
+    // Log warnings for parameters that couldn't be resolved
+    for (const name of names) {
+      if (!parameterMap.has(name)) {
         logger.warn(
-          `Failed to resolve image id from SSM parameter ${invalidParam}: Parameter not found or access denied`,
+          `Failed to resolve image id from SSM parameter ${name}: Parameter not found or access denied`,
         );
-        results.set(invalidParam, undefined);
-      }
-    } catch (error: unknown) {
-      logger.warn(`Failed to resolve image ids from SSM parameters ${chunk.join(', ')}`, { error });
-      // Mark all parameters in this chunk as undefined
-      for (const name of chunk) {
-        results.set(name, undefined);
       }
     }
-  }
 
-  // Return values in the same order as input names
-  return names.map((name) => results.get(name));
+    // Return values in the same order as input names
+    return names.map((name) => parameterMap.get(name));
+  } catch (error: unknown) {
+    logger.warn(`Failed to resolve image ids from SSM parameters ${names.join(', ')}`, { error });
+    // Mark all parameters as undefined on failure
+    return names.map(() => undefined);
+  }
 }
 
 /**
@@ -307,7 +289,7 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
   const wildcardPatterns = options.ssmParameterNames.filter((n) => n.startsWith('*'));
 
   // Batch fetch explicit parameter values in chunks of 10 (AWS API limit)
-  const explicitValuesPromise = resolveSsmParameterValues(explicitNames, ssmClient);
+  const explicitValuesPromise = resolveSsmParameterValues(explicitNames);
 
   // Handle wildcard patterns by first discovering matching parameters, then
   // fetching their values
@@ -332,7 +314,7 @@ async function getAmisReferedInSSM(options: AmiCleanupOptions): Promise<(string
           .map((param) => param.Name)
           .filter((name): name is string => name !== undefined);
 
-        return resolveSsmParameterValues(discoveredNames, ssmClient);
+        return resolveSsmParameterValues(discoveredNames);
       } catch (e) {
         logger.warn('Failed to describe SSM parameters using wildcard patterns', { error: e });
         return [];