Skip to content

Commit 2540ab6

Browse files
dgokcindgokcin
committed
feat(lambda): add ami-updater function for automated ami updates
- implement core ami updater logic with ec2 and ssm integration - add comprehensive unit tests for all components - configure eslint and typescript for code quality - include vitest for testing framework - add aws powertools for logging and metrics - implement dry run mode for safe testing - add error handling and logging throughout
1 parent fae86ea commit 2540ab6

File tree

11 files changed

+684
-0
lines changed

11 files changed

+684
-0
lines changed

‎lambdas/functions/ami-updater/.eslintrc.js‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
module.exports = {
2+
parser: '@typescript-eslint/parser',
3+
extends: [
4+
'eslint:recommended',
5+
'plugin:@typescript-eslint/recommended',
6+
],
7+
parserOptions: {
8+
ecmaVersion: 2020,
9+
sourceType: 'module',
10+
},
11+
env: {
12+
node: true,
13+
es6: true,
14+
},
15+
rules: {
16+
'@typescript-eslint/explicit-function-return-type': 'off',
17+
'@typescript-eslint/no-explicit-any': 'off',
18+
'@typescript-eslint/no-unused-vars': ['error', { argsIgnorePattern: '^_' }],
19+
},
20+
};

‎lambdas/functions/ami-updater/index.ts‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import { EC2Client, DescribeImagesCommand, Image } from '@aws-sdk/client-ec2';
2+
import { SSMClient, GetParameterCommand, PutParameterCommand } from '@aws-sdk/client-ssm';
3+
import { Logger } from '@aws-lambda-powertools/logger';
4+
5+
const logger = new Logger({ serviceName: process.env.POWERTOOLS_SERVICE_NAME });
6+
const DRY_RUN = process.env.DRY_RUN?.toLowerCase() === 'true';
7+
const SSM_PARAMETER_NAME = process.env.SSM_PARAMETER_NAME || '/github-action-runners/latest_ami_id';
8+
const AMI_FILTER = JSON.parse(process.env.AMI_FILTER || '{}');
9+
10+
const ec2Client = new EC2Client({});
11+
const ssmClient = new SSMClient({});
12+
13+
async function getLatestAmi(): Promise<Image> {
14+
try {
15+
const command = new DescribeImagesCommand({
16+
Owners: AMI_FILTER.owners,
17+
Filters: AMI_FILTER.filters,
18+
});
19+
20+
const response = await ec2Client.send(command);
21+
const images = response.Images || [];
22+
23+
if (images.length === 0) {
24+
throw new Error('No matching AMIs found');
25+
}
26+
27+
// Sort by creation date to get the latest
28+
const sortedImages = images.sort((a, b) => {
29+
return (b.CreationDate || '').localeCompare(a.CreationDate || '');
30+
});
31+
32+
return sortedImages[0];
33+
} catch (error) {
34+
logger.error('Error getting latest AMI', { error });
35+
throw error;
36+
}
37+
}
38+
39+
async function getCurrentAmiId(): Promise<string | null> {
40+
try {
41+
const command = new GetParameterCommand({
42+
Name: SSM_PARAMETER_NAME,
43+
});
44+
45+
const response = await ssmClient.send(command);
46+
return response.Parameter?.Value || null;
47+
} catch (error: any) {
48+
if (error.name === 'ParameterNotFound') {
49+
logger.info(`Parameter ${SSM_PARAMETER_NAME} not found`);
50+
return null;
51+
}
52+
logger.error('Error getting current AMI ID from SSM', { error });
53+
throw error;
54+
}
55+
}
56+
57+
async function updateAmiParameter(amiId: string): Promise<{ success: boolean; message: string }> {
58+
try {
59+
const currentAmiId = await getCurrentAmiId();
60+
61+
if (currentAmiId === amiId) {
62+
logger.info('SSM parameter already contains latest AMI ID', { amiId });
63+
return { success: true, message: 'Already using latest AMI' };
64+
}
65+
66+
if (DRY_RUN) {
67+
logger.info('Would update SSM parameter', {
68+
from: currentAmiId,
69+
to: amiId,
70+
dryRun: true,
71+
});
72+
return { success: true, message: 'Would update AMI (Dry Run)' };
73+
}
74+
75+
const command = new PutParameterCommand({
76+
Name: SSM_PARAMETER_NAME,
77+
Value: amiId,
78+
Type: 'String',
79+
Overwrite: true,
80+
});
81+
82+
await ssmClient.send(command);
83+
84+
logger.info('Successfully updated SSM parameter', {
85+
from: currentAmiId,
86+
to: amiId,
87+
});
88+
return { success: true, message: 'Updated successfully' };
89+
} catch (error) {
90+
logger.error('Error updating SSM parameter', { error });
91+
return { success: false, message: `Error: ${error}` };
92+
}
93+
}
94+
95+
export const handler = async (event: any, context: any) => {
96+
logger.info('Starting AMI updater', { dryRun: DRY_RUN });
97+
98+
try {
99+
// Get the latest AMI
100+
const latestAmi = await getLatestAmi();
101+
logger.info('Found latest AMI', { amiId: latestAmi.ImageId });
102+
103+
if (!latestAmi.ImageId) {
104+
throw new Error('Latest AMI ID is undefined');
105+
}
106+
107+
// Update SSM parameter
108+
const { success, message } = await updateAmiParameter(latestAmi.ImageId);
109+
110+
return {
111+
statusCode: success ? 200 : 500,
112+
body: {
113+
dryRun: DRY_RUN,
114+
overallSuccess: success,
115+
latestAmi: latestAmi.ImageId,
116+
message: DRY_RUN ? `[DRY RUN] ${message}` : message,
117+
},
118+
};
119+
} catch (error) {
120+
logger.error('Error in lambda execution', { error });
121+
return {
122+
statusCode: 500,
123+
body: {
124+
dryRun: DRY_RUN,
125+
error: String(error),
126+
overallSuccess: false,
127+
},
128+
};
129+
}
130+
};

‎lambdas/functions/ami-updater/package.json‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
{
2+
"name": "@lambda/ami-updater",
3+
"version": "1.0.0",
4+
"description": "Lambda function to update AMIs in launch templates",
5+
"main": "dist/index.js",
6+
"scripts": {
7+
"build": "tsc",
8+
"test": "vitest run",
9+
"test:watch": "vitest",
10+
"lint": "eslint . --ext .ts",
11+
"lint:fix": "eslint . --ext .ts --fix"
12+
},
13+
"dependencies": {
14+
"@aws-lambda-powertools/logger": "^1.17.0",
15+
"@aws-lambda-powertools/metrics": "^1.17.0",
16+
"@aws-lambda-powertools/tracer": "^1.17.0",
17+
"@aws-sdk/client-ec2": "^3.470.0",
18+
"@types/aws-lambda": "^8.10.130"
19+
},
20+
"devDependencies": {
21+
"@types/node": "^20.10.4",
22+
"@typescript-eslint/eslint-plugin": "^6.13.2",
23+
"@typescript-eslint/parser": "^6.13.2",
24+
"eslint": "^8.55.0",
25+
"typescript": "^5.3.3",
26+
"vitest": "^1.0.2"
27+
}
28+
}

‎lambdas/functions/ami-updater/src/__tests__/ami.test.ts‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import { describe, it, expect, vi, beforeEach } from 'vitest';
2+
import { EC2Client, DescribeImagesCommand, DescribeLaunchTemplatesCommand, DescribeLaunchTemplateVersionsCommand, CreateLaunchTemplateVersionCommand, ModifyLaunchTemplateCommand } from '@aws-sdk/client-ec2';
3+
import { AMIManager } from '../ami';
4+
5+
vi.mock('@aws-sdk/client-ec2');
6+
vi.mock('../../shared/aws-powertools-util', () => ({
7+
logger: {
8+
info: vi.fn(),
9+
warn: vi.fn(),
10+
error: vi.fn(),
11+
},
12+
}));
13+
14+
describe('AMIManager', () => {
15+
let ec2Client: EC2Client;
16+
let amiManager: AMIManager;
17+
18+
beforeEach(() => {
19+
ec2Client = new EC2Client({});
20+
amiManager = new AMIManager(ec2Client);
21+
vi.clearAllMocks();
22+
});
23+
24+
describe('getLatestAmi', () => {
25+
it('should return the latest AMI ID', async () => {
26+
const mockResponse = {
27+
Images: [
28+
{ ImageId: 'ami-2', CreationDate: '2023-12-02' },
29+
{ ImageId: 'ami-1', CreationDate: '2023-12-01' },
30+
],
31+
};
32+
33+
vi.mocked(ec2Client.send).mockResolvedValueOnce(mockResponse);
34+
35+
const config = {
36+
owners: ['self'],
37+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
38+
};
39+
40+
const result = await amiManager.getLatestAmi(config);
41+
expect(result).toBe('ami-2');
42+
expect(ec2Client.send).toHaveBeenCalledWith(expect.any(DescribeImagesCommand));
43+
});
44+
45+
it('should throw error when no AMIs found', async () => {
46+
vi.mocked(ec2Client.send).mockResolvedValueOnce({ Images: [] });
47+
48+
const config = {
49+
owners: ['self'],
50+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
51+
};
52+
53+
await expect(amiManager.getLatestAmi(config)).rejects.toThrow('No matching AMIs found');
54+
});
55+
});
56+
57+
describe('updateLaunchTemplate', () => {
58+
it('should update launch template with new AMI ID', async () => {
59+
vi.mocked(ec2Client.send)
60+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplatesCommand
61+
LaunchTemplates: [{ LatestVersionNumber: 1 }],
62+
})
63+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
64+
LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-old' } }],
65+
})
66+
.mockResolvedValueOnce({ // updateLaunchTemplate - DescribeLaunchTemplatesCommand
67+
LaunchTemplates: [{ LatestVersionNumber: 1 }],
68+
});
69+
70+
const result = await amiManager.updateLaunchTemplate('test-template', 'ami-new', false);
71+
72+
expect(result.success).toBe(true);
73+
expect(result.message).toBe('Updated successfully');
74+
expect(ec2Client.send).toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
75+
expect(ec2Client.send).toHaveBeenCalledWith(expect.any(ModifyLaunchTemplateCommand));
76+
});
77+
78+
it('should not update if AMI ID is the same', async () => {
79+
vi.mocked(ec2Client.send)
80+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplatesCommand
81+
LaunchTemplates: [{ LatestVersionNumber: 1 }],
82+
})
83+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
84+
LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-1' } }],
85+
});
86+
87+
const result = await amiManager.updateLaunchTemplate('test-template', 'ami-1', false);
88+
89+
expect(result.success).toBe(true);
90+
expect(result.message).toBe('Already using latest AMI');
91+
expect(ec2Client.send).not.toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
92+
});
93+
94+
it('should handle dry run mode', async () => {
95+
vi.mocked(ec2Client.send)
96+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplatesCommand
97+
LaunchTemplates: [{ LatestVersionNumber: 1 }],
98+
})
99+
.mockResolvedValueOnce({ // getCurrentAmiId - DescribeLaunchTemplateVersionsCommand
100+
LaunchTemplateVersions: [{ LaunchTemplateData: { ImageId: 'ami-old' } }],
101+
});
102+
103+
const result = await amiManager.updateLaunchTemplate('test-template', 'ami-new', true);
104+
105+
expect(result.success).toBe(true);
106+
expect(result.message).toBe('Would update AMI (Dry Run)');
107+
expect(ec2Client.send).not.toHaveBeenCalledWith(expect.any(CreateLaunchTemplateVersionCommand));
108+
});
109+
});
110+
});

‎lambdas/functions/ami-updater/src/__tests__/config.test.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
import { describe, it, expect, vi, beforeEach } from 'vitest';
2+
import { getConfig } from '../config';
3+
4+
vi.mock('../../shared/aws-powertools-util', () => ({
5+
logger: {
6+
error: vi.fn(),
7+
},
8+
}));
9+
10+
describe('getConfig', () => {
11+
beforeEach(() => {
12+
vi.resetModules();
13+
process.env = {};
14+
});
15+
16+
it('should return valid configuration when all environment variables are set correctly', () => {
17+
process.env.LAUNCH_TEMPLATE_NAME = 'test-template';
18+
process.env.DRY_RUN = 'true';
19+
process.env.AMI_FILTER = JSON.stringify({
20+
owners: ['self'],
21+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
22+
});
23+
24+
const config = getConfig();
25+
26+
expect(config).toEqual({
27+
launchTemplateName: 'test-template',
28+
dryRun: true,
29+
amiFilter: {
30+
owners: ['self'],
31+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
32+
},
33+
});
34+
});
35+
36+
it('should handle DRY_RUN=false correctly', () => {
37+
process.env.LAUNCH_TEMPLATE_NAME = 'test-template';
38+
process.env.DRY_RUN = 'false';
39+
process.env.AMI_FILTER = JSON.stringify({
40+
owners: ['self'],
41+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
42+
});
43+
44+
const config = getConfig();
45+
expect(config.dryRun).toBe(false);
46+
});
47+
48+
it('should throw error when LAUNCH_TEMPLATE_NAME is not set', () => {
49+
process.env.AMI_FILTER = JSON.stringify({
50+
owners: ['self'],
51+
filters: [{ name: 'tag:Environment', values: ['prod'] }],
52+
});
53+
54+
expect(() => getConfig()).toThrow('LAUNCH_TEMPLATE_NAME environment variable is not set');
55+
});
56+
57+
it('should throw error when AMI_FILTER is not set', () => {
58+
process.env.LAUNCH_TEMPLATE_NAME = 'test-template';
59+
60+
expect(() => getConfig()).toThrow('AMI_FILTER environment variable is not set');
61+
});
62+
63+
it('should throw error when AMI_FILTER is invalid JSON', () => {
64+
process.env.LAUNCH_TEMPLATE_NAME = 'test-template';
65+
process.env.AMI_FILTER = 'invalid-json';
66+
67+
expect(() => getConfig()).toThrow('Invalid AMI_FILTER format');
68+
});
69+
70+
it('should throw error when AMI_FILTER has invalid structure', () => {
71+
process.env.LAUNCH_TEMPLATE_NAME = 'test-template';
72+
process.env.AMI_FILTER = JSON.stringify({
73+
owners: 'not-an-array',
74+
filters: 'not-an-array',
75+
});
76+
77+
expect(() => getConfig()).toThrow('AMI_FILTER must contain owners (array) and filters (array)');
78+
});
79+
});

0 commit comments

Comments
 (0)