From dfba3e20f81cbb6c418b0aeef423a3c6b30fc1a9 Mon Sep 17 00:00:00 2001
From: Luca Pizzini <lpizzini7@gmail.com>
Date: Wed, 7 Jun 2023 09:51:08 +0200
Subject: [PATCH] fix(codebuild): add possibility to specify
 BUILD_GENERAL1_SMALL compute type with Linux GPU build image

---
 .../lib/linux-gpu-build-image.ts              |  5 +--
 .../test/linux-gpu-build-image.test.ts        | 33 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts b/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
index a31678a5a6567..218703089b643 100644
--- a/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
+++ b/packages/aws-cdk-lib/aws-codebuild/lib/linux-gpu-build-image.ts
@@ -142,8 +142,9 @@ export class LinuxGpuBuildImage implements IBindableBuildImage {
   public validate(buildEnvironment: BuildEnvironment): string[] {
     const ret = [];
     if (buildEnvironment.computeType &&
-        buildEnvironment.computeType !== ComputeType.LARGE) {
-      ret.push(`GPU images only support ComputeType '${ComputeType.LARGE}' - ` +
+        buildEnvironment.computeType !== ComputeType.LARGE &&
+        buildEnvironment.computeType !== ComputeType.SMALL) {
+      ret.push(`GPU images only support ComputeType '${ComputeType.LARGE}' and '${ComputeType.SMALL}' - ` +
         `'${buildEnvironment.computeType}' was given`);
     }
     return ret;
diff --git a/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts b/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
index f1fdada6c68e5..e090c778f9bae 100644
--- a/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
+++ b/packages/aws-cdk-lib/aws-codebuild/test/linux-gpu-build-image.test.ts
@@ -57,6 +57,39 @@ describe('Linux GPU build image', () => {
         },
       });
     });
+
+    test('allows to specify a BUILD_GENERAL1_SMALL build environment', () => {
+      const stack = new cdk.Stack();
+
+      new codebuild.Project(stack, 'Project', {
+        buildSpec: codebuild.BuildSpec.fromObject({
+          version: '0.2',
+          phases: {
+            build: { commands: ['ls'] },
+          },
+        }),
+        environment: {
+          buildImage: codebuild.LinuxGpuBuildImage.awsDeepLearningContainersImage(
+            'my-repo', 'my-tag', '123456789012'),
+          computeType: codebuild.ComputeType.SMALL,
+        },
+      });
+
+      Template.fromStack(stack).hasResourceProperties('AWS::CodeBuild::Project', {
+        Environment: {
+          ComputeType: 'BUILD_GENERAL1_SMALL',
+          Image: {
+            'Fn::Join': ['', [
+              '123456789012.dkr.ecr.',
+              { Ref: 'AWS::Region' },
+              '.',
+              { Ref: 'AWS::URLSuffix' },
+              '/my-repo:my-tag',
+            ]],
+          },
+        },
+      });
+    });
   });
 
   describe('ECR Repository', () => {