From 7331e328234348453547f80f4253485fb3ebc2b0 Mon Sep 17 00:00:00 2001
From: go-to-k <24818752+go-to-k@users.noreply.github.com>
Date: Sat, 30 Dec 2023 02:06:20 +0900
Subject: [PATCH] feat(stepfunctions-tasks): add new allocation strategies for
 spot instance fleets in EmrCreateCluster

change unit tests

change unit tests
---
 .../lib/emr/emr-create-cluster.ts                  | 14 ++++++++++++++
 .../test/emr/emr-create-cluster.test.ts            | 13 +++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
index a1f82782755c4..804c1577b187a 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
@@ -680,6 +680,20 @@ export namespace EmrCreateCluster {
      * Capacity-optimized, which launches instances from Spot Instance pools with optimal capacity for the number of instances that are launching.
      */
     CAPACITY_OPTIMIZED = 'capacity-optimized',
+    /**
+     * Price-capacity-optimized, which launches instances from Spot Instance pools with the highest capacity availability for the number of instances that are launching.
+     *
+     * Recommended.
+     */
+    PRICE_CAPACITY_OPTIMIZED = 'price-capacity-optimized',
+    /**
+     * Lowest-price, which launches instances from the lowest priced pool that has available capacity.
+     */
+    LOWEST_PRICE = 'lowest-price',
+    /**
+     * Diversified, which launches instances across all Spot capacity pools.
+     */
+    DIVERSIFIED = 'diversified',
   }
 
   /**
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
index a0be73f8f94ba..8274f08232386 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
@@ -1,4 +1,4 @@
-import { Template, Match } from '../../../assertions';
+import { Template } from '../../../assertions';
 import * as iam from '../../../aws-iam';
 import * as sfn from '../../../aws-stepfunctions';
 import * as cdk from '../../../core';
@@ -883,7 +883,12 @@ test('Create Cluster with Instances configuration', () => {
   });
 });
 
-test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimized for Spot instances', () => {
+test.each([
+  [EmrCreateCluster.SpotAllocationStrategy.CAPACITY_OPTIMIZED, 'capacity-optimized'],
+  [EmrCreateCluster.SpotAllocationStrategy.PRICE_CAPACITY_OPTIMIZED, 'price-capacity-optimized'],
+  [EmrCreateCluster.SpotAllocationStrategy.LOWEST_PRICE, 'lowest-price'],
+  [EmrCreateCluster.SpotAllocationStrategy.DIVERSIFIED, 'diversified'],
+])('Create Cluster with InstanceFleet with allocation strategy %s for Spot instances', (strategy, expected) => {
   // WHEN
   const task = new EmrCreateCluster(stack, 'Task', {
     instances: {
@@ -914,7 +919,7 @@ test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimi
         }],
         launchSpecifications: {
           spotSpecification: {
-            allocationStrategy: EmrCreateCluster.SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+            allocationStrategy: strategy,
             blockDurationMinutes: 1,
             timeoutAction: EmrCreateCluster.SpotTimeoutAction.TERMINATE_CLUSTER,
             timeoutDurationMinutes: 1,
@@ -978,7 +983,7 @@ test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimi
           }],
           LaunchSpecifications: {
             SpotSpecification: {
-              AllocationStrategy: 'capacity-optimized',
+              AllocationStrategy: expected,
               BlockDurationMinutes: 1,
               TimeoutAction: 'TERMINATE_CLUSTER',
               TimeoutDurationMinutes: 1,