add unit tests for isGpuInstanceType

pahud · Oct 6, 2024 · 8fd3db8 · 8fd3db8
1 parent 8e88f82
commit 8fd3db8
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 12 deletions.
diff --git a/packages/aws-cdk-lib/aws-eks/lib/managed-nodegroup.ts b/packages/aws-cdk-lib/aws-eks/lib/managed-nodegroup.ts
@@ -5,6 +5,7 @@ import { InstanceType, ISecurityGroup, SubnetSelection, InstanceArchitecture, In
 import { IRole, ManagedPolicy, PolicyStatement, Role, ServicePrincipal } from '../../aws-iam';
 import { IResource, Resource, Annotations, withResolved, FeatureFlags } from '../../core';
 import * as cxapi from '../../cx-api';
+import { isGpuInstanceType } from './private/nodegroup';
 
 /**
  * NodeGroup interface
@@ -598,18 +599,6 @@ const gpuAmiTypes: NodegroupAmiType[] = [
   NodegroupAmiType.BOTTLEROCKET_ARM_64_NVIDIA,
 ];
 
-/**
- * This function check if the instanceType is GPU instance.
- * @param instanceType The EC2 instance type
- */
-function isGpuInstanceType(instanceType: InstanceType): boolean {
-  //compare instanceType to known GPU InstanceTypes
-  const knownGpuInstanceTypes = [InstanceClass.P2, InstanceClass.P3, InstanceClass.P3DN, InstanceClass.P4DE, InstanceClass.P4D,
-    InstanceClass.G3S, InstanceClass.G3, InstanceClass.G4DN, InstanceClass.G4AD, InstanceClass.G5, InstanceClass.G5G, InstanceClass.G6,
-    InstanceClass.G6E, InstanceClass.INF1, InstanceClass.INF2];
-  return knownGpuInstanceTypes.some((c) => instanceType.sameInstanceClassAs(InstanceType.of(c, InstanceSize.LARGE)));
-}
-
 /**
  * This function check if the instanceType is supported by Windows AMI.
  * https://docs.aws.amazon.com/eks/latest/userguide/windows-support.html

diff --git a/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts b/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts
@@ -0,0 +1,13 @@
+import { InstanceClass, InstanceSize, InstanceType } from '../../../aws-ec2';
+
+/**
+ * This function check if the instanceType is GPU instance.
+ * @param instanceType The EC2 instance type
+ */
+export function isGpuInstanceType(instanceType: InstanceType): boolean {
+  //compare instanceType to known GPU InstanceTypes
+  const knownGpuInstanceTypes = [InstanceClass.P2, InstanceClass.P3, InstanceClass.P3DN, InstanceClass.P4DE, InstanceClass.P4D,
+    InstanceClass.G3S, InstanceClass.G3, InstanceClass.G4DN, InstanceClass.G4AD, InstanceClass.G5, InstanceClass.G5G, InstanceClass.G6,
+    InstanceClass.G6E, InstanceClass.INF1, InstanceClass.INF2];
+  return knownGpuInstanceTypes.some((c) => instanceType.sameInstanceClassAs(InstanceType.of(c, InstanceSize.LARGE)));
+}
diff --git a/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts b/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts
@@ -7,6 +7,7 @@ import * as cdk from '../../core';
 import * as cxapi from '../../cx-api';
 import * as eks from '../lib';
 import { NodegroupAmiType, TaintEffect } from '../lib';
+import { isGpuInstanceType } from '../lib/private/nodegroup';
 
 /* eslint-disable max-len */
 
@@ -1735,3 +1736,42 @@ describe('node group', () => {
     expect(() => cluster.addNodegroupCapacity('ng', { maxUnavailablePercentage: 101 })).toThrow(/maxUnavailablePercentage must be between 1 and 100/);
   });
 });
+
+describe('isGpuec2.InstanceType', () => {
+  it('should return true for known GPU instance types', () => {
+    const gpuInstances = [
+      ec2.InstanceType.of(ec2.InstanceClass.P2, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.G3, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.P4D, ec2.InstanceSize.LARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.G6, ec2.InstanceSize.MEDIUM),
+    ];
+
+    gpuInstances.forEach(instance => {
+      expect(isGpuInstanceType(instance)).toBe(true);
+    });
+  });
+
+  it('should return false for non-GPU instance types', () => {
+    const nonGpuInstances = [
+      ec2.InstanceType.of(ec2.InstanceClass.T3, ec2.InstanceSize.MICRO),
+      ec2.InstanceType.of(ec2.InstanceClass.M5, ec2.InstanceSize.LARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.C5, ec2.InstanceSize.XLARGE),
+    ];
+
+    nonGpuInstances.forEach(instance => {
+      expect(isGpuInstanceType(instance)).toBe(false);
+    });
+  });
+
+  it('should return true for different sizes of GPU instance types', () => {
+    const gpuInstances = [
+      ec2.InstanceType.of(ec2.InstanceClass.P3, ec2.InstanceSize.LARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.P3, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.P3, ec2.InstanceSize.METAL),
+    ];
+
+    gpuInstances.forEach(instance => {
+      expect(isGpuInstanceType(instance)).toBe(true);
+    });
+  });
+});