From 8cc345d86dd89ce1718e50278636edcfc8eb295c Mon Sep 17 00:00:00 2001
From: Oguzhan Unlu <oguzhan@snowplowanalytics.com>
Date: Mon, 18 Mar 2024 17:22:26 +0300
Subject: [PATCH] Add gp3 volume support (close #93)

---
 src/emr_cluster.go      |  2 +-
 src/emr_cluster_test.go | 69 +++++++++++++++++++++++++++++++++++
 src/test_data_test.go   | 81 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 151 insertions(+), 1 deletion(-)

diff --git a/src/emr_cluster.go b/src/emr_cluster.go
index 9fc2610..ea77bce 100644
--- a/src/emr_cluster.go
+++ b/src/emr_cluster.go
@@ -310,7 +310,7 @@ func GetEbsConfiguration(c *EbsConfigurationRecord) *emr.EbsConfiguration {
 			emrVolumeSpec := &emr.VolumeSpecification{}
 			emrVolumeSpec.SizeInGB = aws.Int64(config.VolumeSpecification.SizeInGB)
 			emrVolumeSpec.VolumeType = aws.String(config.VolumeSpecification.VolumeType)
-			if *emrVolumeSpec.VolumeType != "gp2" {
+			if *emrVolumeSpec.VolumeType != "gp2" && *emrVolumeSpec.VolumeType != "gp3" {
 				emrVolumeSpec.Iops = aws.Int64(config.VolumeSpecification.Iops)
 			}
 
diff --git a/src/emr_cluster_test.go b/src/emr_cluster_test.go
index 1297c8d..a839283 100644
--- a/src/emr_cluster_test.go
+++ b/src/emr_cluster_test.go
@@ -395,6 +395,75 @@ func TestGetInstanceGroups_WithEBS(t *testing.T) {
 	assert.Equal(expected[2], groups[2])
 }
 
+func TestGetInstanceGroups_WithGP3(t *testing.T) {
+	assert := assert.New(t)
+
+	record, _ := CR.ParseClusterRecord([]byte(ClusterRecordWithGP3), nil, "")
+	ec, _ := InitEmrCluster(*record)
+	groups := ec.GetInstanceGroups()
+	assert.Len(groups, 3)
+	expected := []*emr.InstanceGroupConfig{
+		{
+			InstanceCount: aws.Int64(1),
+			InstanceRole:  aws.String("MASTER"),
+			InstanceType:  aws.String("m1.medium"),
+			EbsConfiguration: &emr.EbsConfiguration{
+				EbsOptimized: aws.Bool(true),
+				EbsBlockDeviceConfigs: []*emr.EbsBlockDeviceConfig{
+					{
+						VolumesPerInstance: aws.Int64(12),
+						VolumeSpecification: &emr.VolumeSpecification{
+							SizeInGB:   aws.Int64(10),
+							VolumeType: aws.String("gp3"),
+						},
+					},
+				},
+			},
+		},
+		{
+			InstanceCount: aws.Int64(3),
+			InstanceRole:  aws.String("CORE"),
+			InstanceType:  aws.String("c3.4xlarge"),
+			EbsConfiguration: &emr.EbsConfiguration{
+				EbsOptimized: aws.Bool(false),
+				EbsBlockDeviceConfigs: []*emr.EbsBlockDeviceConfig{
+					{
+						VolumesPerInstance: aws.Int64(8),
+						VolumeSpecification: &emr.VolumeSpecification{
+							Iops:       aws.Int64(20),
+							SizeInGB:   aws.Int64(4),
+							VolumeType: aws.String("io1"),
+						},
+					},
+				},
+			},
+		},
+		{
+			InstanceCount: aws.Int64(1),
+			InstanceRole:  aws.String("TASK"),
+			InstanceType:  aws.String("m1.medium"),
+			BidPrice:      aws.String("0.015"),
+			Market:        aws.String("SPOT"),
+			EbsConfiguration: &emr.EbsConfiguration{
+				EbsOptimized: aws.Bool(false),
+				EbsBlockDeviceConfigs: []*emr.EbsBlockDeviceConfig{
+					{
+						VolumesPerInstance: aws.Int64(4),
+						VolumeSpecification: &emr.VolumeSpecification{
+							Iops:       aws.Int64(100),
+							SizeInGB:   aws.Int64(6),
+							VolumeType: aws.String("standard"),
+						},
+					},
+				},
+			},
+		},
+	}
+	assert.Equal(expected[0], groups[0])
+	assert.Equal(expected[1], groups[1])
+	assert.Equal(expected[2], groups[2])
+}
+
 func TestGetTags_NoTags(t *testing.T) {
 	record, _ := CR.ParseClusterRecord([]byte(ClusterRecord1), nil, "")
 	ec, _ := InitEmrCluster(*record)
diff --git a/src/test_data_test.go b/src/test_data_test.go
index cc4c122..e1c7c85 100644
--- a/src/test_data_test.go
+++ b/src/test_data_test.go
@@ -177,6 +177,87 @@ var ClusterRecordWithEBS = `{
   }
 }`
 
+var ClusterRecordWithGP3 = `{
+  "schema": "iglu:com.snowplowanalytics.dataflowrunner/ClusterConfig/avro/1-0-0",
+  "data": {
+    "name": "xxx",
+    "logUri": "s3://logging/",
+    "region": "us-east-1",
+    "credentials": {
+      "accessKeyId": "env",
+      "secretAccessKey": "env"
+    },
+    "roles": {
+      "jobflow": "EMR_EC2_DefaultRole",
+      "service": "EMR_DefaultRole"
+    },
+    "ec2": {
+      "amiVersion": "4.5.0",
+      "keyName": "snowplow-yyy-key",
+      "location": {
+        "classic": {
+          "availabilityZone": "us-east-1a"
+        },
+        "vpc": {
+          "subnetId": "subnet-123456"
+        }
+      },
+      "instances": {
+        "master": {
+          "type": "m1.medium",
+          "ebsConfiguration": {
+            "ebsOptimized": true,
+            "ebsBlockDeviceConfigs": [
+              {
+                "volumesPerInstance": 12,
+                "volumeSpecification": {
+                  "sizeInGB": 10,
+                  "volumeType": "gp3"
+                }
+              }
+            ]
+          }
+        },
+        "core": {
+          "type": "c3.4xlarge",
+          "count": 3,
+          "ebsConfiguration": {
+            "ebsOptimized": false,
+            "ebsBlockDeviceConfigs": [
+              {
+                "volumesPerInstance": 8,
+                "volumeSpecification": {
+                  "iops": 20,
+                  "sizeInGB": 4,
+                  "volumeType": "io1"
+                }
+              }
+            ]
+          }
+        },
+        "task": {
+          "type": "m1.medium",
+          "count": 1,
+          "bid": "0.015",
+          "ebsConfiguration": {
+            "ebsOptimized": false,
+            "ebsBlockDeviceConfigs": [
+              {
+                "volumesPerInstance": 4,
+                "volumeSpecification": {
+                  "iops": 100,
+                  "sizeInGB": 6,
+                  "volumeType": "standard"
+                }
+              }
+            ]
+          }
+        }
+      }
+    }
+  }
+}`
+
 var ClusterRecordWithApps = `{
   "schema": "iglu:com.snowplowanalytics.dataflowrunner/ClusterConfig/avro/1-0-0",
   "data": {