@@ -7,19 +7,17 @@ import {
77 EC2Client ,
88 Image ,
99} from '@aws-sdk/client-ec2' ;
10- import {
11- DescribeParametersCommand ,
12- DescribeParametersCommandOutput ,
13- GetParameterCommand ,
14- SSMClient ,
15- } from '@aws-sdk/client-ssm' ;
10+ import { DescribeParametersCommand , DescribeParametersCommandOutput , SSMClient } from '@aws-sdk/client-ssm' ;
11+ import { getParameters } from '@aws-github-runner/aws-ssm-util' ;
1612import { mockClient } from 'aws-sdk-client-mock' ;
1713import 'aws-sdk-client-mock-jest/vitest' ;
1814
1915import { AmiCleanupOptions , amiCleanup , defaultAmiCleanupOptions } from './ami' ;
2016import { describe , it , expect , beforeEach , vi } from 'vitest' ;
2117import { fail } from 'assert' ;
2218
19+ vi . mock ( '@aws-github-runner/aws-ssm-util' ) ;
20+
2321process . env . AWS_REGION = 'eu-east-1' ;
2422const deleteAmisOlderThenDays = 30 ;
2523const date31DaysAgo = new Date ( new Date ( ) . setDate ( new Date ( ) . getDate ( ) - ( deleteAmisOlderThenDays + 1 ) ) ) ;
@@ -83,22 +81,12 @@ describe("delete AMI's", () => {
8381 mockSSMClient . reset ( ) ;
8482
8583 mockSSMClient . on ( DescribeParametersCommand ) . resolves ( ssmParameters ) ;
86- mockSSMClient . on ( GetParameterCommand , { Name : 'ami-id/ami-ssm0001' } ) . resolves ( {
87- Parameter : {
88- Name : 'ami-id/ami-ssm0001' ,
89- Type : 'String' ,
90- Value : 'ami-ssm0001' ,
91- Version : 1 ,
92- } ,
93- } ) ;
94- mockSSMClient . on ( GetParameterCommand , { Name : 'ami-id/ami-ssm0002' } ) . resolves ( {
95- Parameter : {
96- Name : 'ami-id/ami-ssm0002' ,
97- Type : 'String' ,
98- Value : 'ami-ssm0002' ,
99- Version : 1 ,
100- } ,
101- } ) ;
84+ vi . mocked ( getParameters ) . mockResolvedValue (
85+ new Map ( [
86+ [ 'ami-id/ami-ssm0001' , 'ami-ssm0001' ] ,
87+ [ 'ami-id/ami-ssm0002' , 'ami-ssm0002' ] ,
88+ ] ) ,
89+ ) ;
10290
10391 mockEC2Client . on ( DescribeLaunchTemplatesCommand ) . resolves ( {
10492 LaunchTemplates : [
@@ -143,13 +131,7 @@ describe("delete AMI's", () => {
143131 expect ( mockEC2Client ) . toHaveReceivedCommand ( DescribeLaunchTemplatesCommand ) ;
144132 expect ( mockEC2Client ) . toHaveReceivedCommand ( DescribeLaunchTemplateVersionsCommand ) ;
145133 expect ( mockSSMClient ) . toHaveReceivedCommand ( DescribeParametersCommand ) ;
146- expect ( mockSSMClient ) . toHaveReceivedCommandTimes ( GetParameterCommand , 2 ) ;
147- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
148- Name : 'ami-id/ami-ssm0001' ,
149- } ) ;
150- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
151- Name : 'ami-id/ami-ssm0002' ,
152- } ) ;
134+ expect ( getParameters ) . toHaveBeenCalledWith ( [ 'ami-id/ami-ssm0001' , 'ami-id/ami-ssm0002' ] ) ;
153135 } ) ;
154136
155137 it ( 'should NOT delete instances in use.' , async ( ) => {
@@ -485,14 +467,7 @@ describe("delete AMI's", () => {
485467 ] ,
486468 } ) ;
487469
488- mockSSMClient . on ( GetParameterCommand , { Name : '/github-runner/config/ami_id' } ) . resolves ( {
489- Parameter : {
490- Name : '/github-runner/config/ami_id' ,
491- Type : 'String' ,
492- Value : 'ami-underscore0001' ,
493- Version : 1 ,
494- } ,
495- } ) ;
470+ vi . mocked ( getParameters ) . mockResolvedValue ( new Map ( [ [ '/github-runner/config/ami_id' , 'ami-underscore0001' ] ] ) ) ;
496471
497472 await amiCleanup ( {
498473 minimumDaysOld : 0 ,
@@ -501,9 +476,7 @@ describe("delete AMI's", () => {
501476
502477 // AMI should not be deleted because it's referenced in SSM
503478 expect ( mockEC2Client ) . not . toHaveReceivedCommand ( DeregisterImageCommand ) ;
504- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
505- Name : '/github-runner/config/ami_id' ,
506- } ) ;
479+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/github-runner/config/ami_id' ] ) ;
507480 expect ( mockSSMClient ) . not . toHaveReceivedCommand ( DescribeParametersCommand ) ;
508481 } ) ;
509482
@@ -518,14 +491,7 @@ describe("delete AMI's", () => {
518491 ] ,
519492 } ) ;
520493
521- mockSSMClient . on ( GetParameterCommand , { Name : '/github-runner/config/ami-id' } ) . resolves ( {
522- Parameter : {
523- Name : '/github-runner/config/ami-id' ,
524- Type : 'String' ,
525- Value : 'ami-hyphen0001' ,
526- Version : 1 ,
527- } ,
528- } ) ;
494+ vi . mocked ( getParameters ) . mockResolvedValue ( new Map ( [ [ '/github-runner/config/ami-id' , 'ami-hyphen0001' ] ] ) ) ;
529495
530496 await amiCleanup ( {
531497 minimumDaysOld : 0 ,
@@ -534,9 +500,7 @@ describe("delete AMI's", () => {
534500
535501 // AMI should not be deleted because it's referenced in SSM
536502 expect ( mockEC2Client ) . not . toHaveReceivedCommand ( DeregisterImageCommand ) ;
537- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
538- Name : '/github-runner/config/ami-id' ,
539- } ) ;
503+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/github-runner/config/ami-id' ] ) ;
540504 expect ( mockSSMClient ) . not . toHaveReceivedCommand ( DescribeParametersCommand ) ;
541505 } ) ;
542506
@@ -561,14 +525,7 @@ describe("delete AMI's", () => {
561525 ] ,
562526 } ) ;
563527
564- mockSSMClient . on ( GetParameterCommand , { Name : '/some/path/ami-id' } ) . resolves ( {
565- Parameter : {
566- Name : '/some/path/ami-id' ,
567- Type : 'String' ,
568- Value : 'ami-wildcard0001' ,
569- Version : 1 ,
570- } ,
571- } ) ;
528+ vi . mocked ( getParameters ) . mockResolvedValue ( new Map ( [ [ '/some/path/ami-id' , 'ami-wildcard0001' ] ] ) ) ;
572529
573530 await amiCleanup ( {
574531 minimumDaysOld : 0 ,
@@ -580,9 +537,7 @@ describe("delete AMI's", () => {
580537 expect ( mockSSMClient ) . toHaveReceivedCommandWith ( DescribeParametersCommand , {
581538 ParameterFilters : [ { Key : 'Name' , Option : 'Contains' , Values : [ 'ami-id' ] } ] ,
582539 } ) ;
583- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
584- Name : '/some/path/ami-id' ,
585- } ) ;
540+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/some/path/ami-id' ] ) ;
586541 } ) ;
587542
588543 it ( 'handles wildcard SSM parameter patterns (*ami_id)' , async ( ) => {
@@ -606,14 +561,9 @@ describe("delete AMI's", () => {
606561 ] ,
607562 } ) ;
608563
609- mockSSMClient . on ( GetParameterCommand , { Name : '/github-runner/config/ami_id' } ) . resolves ( {
610- Parameter : {
611- Name : '/github-runner/config/ami_id' ,
612- Type : 'String' ,
613- Value : 'ami-wildcard-underscore0001' ,
614- Version : 1 ,
615- } ,
616- } ) ;
564+ vi . mocked ( getParameters ) . mockResolvedValue (
565+ new Map ( [ [ '/github-runner/config/ami_id' , 'ami-wildcard-underscore0001' ] ] ) ,
566+ ) ;
617567
618568 await amiCleanup ( {
619569 minimumDaysOld : 0 ,
@@ -625,9 +575,7 @@ describe("delete AMI's", () => {
625575 expect ( mockSSMClient ) . toHaveReceivedCommandWith ( DescribeParametersCommand , {
626576 ParameterFilters : [ { Key : 'Name' , Option : 'Contains' , Values : [ 'ami_id' ] } ] ,
627577 } ) ;
628- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
629- Name : '/github-runner/config/ami_id' ,
630- } ) ;
578+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/github-runner/config/ami_id' ] ) ;
631579 } ) ;
632580
633581 it ( 'handles mixed explicit names and wildcard patterns' , async ( ) => {
@@ -649,14 +597,9 @@ describe("delete AMI's", () => {
649597 ] ,
650598 } ) ;
651599
652- mockSSMClient . on ( GetParameterCommand , { Name : '/explicit/ami_id' } ) . resolves ( {
653- Parameter : {
654- Name : '/explicit/ami_id' ,
655- Type : 'String' ,
656- Value : 'ami-explicit0001' ,
657- Version : 1 ,
658- } ,
659- } ) ;
600+ vi . mocked ( getParameters )
601+ . mockResolvedValueOnce ( new Map ( [ [ '/explicit/ami_id' , 'ami-explicit0001' ] ] ) )
602+ . mockResolvedValueOnce ( new Map ( [ [ '/discovered/ami-id' , 'ami-wildcard0001' ] ] ) ) ;
660603
661604 mockSSMClient . on ( DescribeParametersCommand ) . resolves ( {
662605 Parameters : [
@@ -668,15 +611,6 @@ describe("delete AMI's", () => {
668611 ] ,
669612 } ) ;
670613
671- mockSSMClient . on ( GetParameterCommand , { Name : '/discovered/ami-id' } ) . resolves ( {
672- Parameter : {
673- Name : '/discovered/ami-id' ,
674- Type : 'String' ,
675- Value : 'ami-wildcard0001' ,
676- Version : 1 ,
677- } ,
678- } ) ;
679-
680614 await amiCleanup ( {
681615 minimumDaysOld : 0 ,
682616 ssmParameterNames : [ '/explicit/ami_id' , '*ami-id' ] ,
@@ -688,15 +622,11 @@ describe("delete AMI's", () => {
688622 ImageId : 'ami-unused0001' ,
689623 } ) ;
690624
691- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
692- Name : '/explicit/ami_id' ,
693- } ) ;
625+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/explicit/ami_id' ] ) ;
694626 expect ( mockSSMClient ) . toHaveReceivedCommandWith ( DescribeParametersCommand , {
695627 ParameterFilters : [ { Key : 'Name' , Option : 'Contains' , Values : [ 'ami-id' ] } ] ,
696628 } ) ;
697- expect ( mockSSMClient ) . toHaveReceivedCommandWith ( GetParameterCommand , {
698- Name : '/discovered/ami-id' ,
699- } ) ;
629+ expect ( getParameters ) . toHaveBeenCalledWith ( [ '/discovered/ami-id' ] ) ;
700630 } ) ;
701631
702632 it ( 'handles SSM parameter fetch failures gracefully' , async ( ) => {
@@ -710,7 +640,7 @@ describe("delete AMI's", () => {
710640 ] ,
711641 } ) ;
712642
713- mockSSMClient . on ( GetParameterCommand , { Name : '/nonexistent/ami_id' } ) . rejects ( new Error ( 'ParameterNotFound' ) ) ;
643+ vi . mocked ( getParameters ) . mockRejectedValue ( new Error ( 'ParameterNotFound' ) ) ;
714644
715645 // Should not throw and should delete the AMI since SSM reference failed
716646 await amiCleanup ( {
@@ -768,7 +698,7 @@ describe("delete AMI's", () => {
768698 ImageId : 'ami-no-ssm0001' ,
769699 } ) ;
770700 expect ( mockSSMClient ) . not . toHaveReceivedCommand ( DescribeParametersCommand ) ;
771- expect ( mockSSMClient ) . not . toHaveReceivedCommand ( GetParameterCommand ) ;
701+ expect ( getParameters ) . not . toHaveBeenCalled ( ) ;
772702 } ) ;
773703 } ) ;
774704} ) ;
0 commit comments