From 288c467499c9671173231e7188e79703504018b8 Mon Sep 17 00:00:00 2001
From: Indresh-Prakash <indreshprakash24@gmail.com>
Date: Sun, 2 Jul 2023 15:33:13 +0530
Subject: [PATCH] Remove check for Throughput

Signed-off-by: Indresh-Prakash <indreshprakash24@gmail.com>
---
 pkg/cloud/cloud.go      |  2 +-
 pkg/cloud/cloud_test.go | 58 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/pkg/cloud/cloud.go b/pkg/cloud/cloud.go
index 70ee16a740..a6798369b0 100644
--- a/pkg/cloud/cloud.go
+++ b/pkg/cloud/cloud.go
@@ -479,7 +479,7 @@ func (c *cloud) ModifyDisk(ctx context.Context, volumeID string, options *Modify
 	if options.VolumeType != "" {
 		req.VolumeType = aws.String(options.VolumeType)
 	}
-	if options.VolumeType == VolumeTypeGP3 {
+	if options.Throughput != 0 {
 		req.Throughput = aws.Int64(int64(options.Throughput))
 	}
 
diff --git a/pkg/cloud/cloud_test.go b/pkg/cloud/cloud_test.go
index 99860fced1..439d6f3472 100644
--- a/pkg/cloud/cloud_test.go
+++ b/pkg/cloud/cloud_test.go
@@ -710,6 +710,64 @@ func TestDeleteDisk(t *testing.T) {
 	}
 }
 
+func TestModifyDisk(t *testing.T) {
+	testCases := []struct {
+		name                string
+		volumeID            string
+		modifyDiskOptions   *ModifyDiskOptions
+		modifiedVolume      *ec2.ModifyVolumeOutput
+		modifiedVolumeError awserr.Error
+	}{
+		{
+			name:     "success: IOPS with GP3",
+			volumeID: "vol-test-1234",
+			modifyDiskOptions: &ModifyDiskOptions{
+				VolumeType: "GP3",
+				IOPS:       3000,
+			},
+			modifiedVolume: &ec2.ModifyVolumeOutput{
+				VolumeModification: &ec2.VolumeModification{
+					VolumeId:          aws.String("vol-test"),
+					TargetIops:        aws.Int64(3000),
+					ModificationState: aws.String(ec2.VolumeModificationStateCompleted),
+				},
+			},
+		},
+		{
+			name:                "fail: ModifyVolume returned generic error",
+			volumeID:            "vol-test-1234",
+			modifiedVolumeError: awserr.New("InvalidParameterCombination", "The parameter iops is not supported for gp2 volumes", nil),
+			modifyDiskOptions: &ModifyDiskOptions{
+				VolumeType: "GP2",
+				IOPS:       3000,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mockCtrl := gomock.NewController(t)
+			mockEC2 := NewMockEC2API(mockCtrl)
+			c := newCloud(mockEC2)
+			ctx := context.Background()
+			if tc.modifiedVolumeError != nil {
+				mockEC2.EXPECT().ModifyVolumeWithContext(gomock.Any(), gomock.Any()).Return(nil, tc.modifiedVolumeError).AnyTimes()
+			} else {
+				mockEC2.EXPECT().ModifyVolumeWithContext(gomock.Any(), gomock.Any()).Return(tc.modifiedVolume, nil).AnyTimes()
+			}
+			err := c.ModifyDisk(ctx, tc.volumeID, tc.modifyDiskOptions)
+			if err != nil && tc.modifiedVolumeError == nil {
+				t.Fatalf("ModifyDisk() failed: expected no error, got: %v", err)
+			}
+
+			if err == nil && tc.modifiedVolumeError != nil {
+				t.Fatal("ModifyDisk() failed: expected error, got nothing")
+			}
+			mockCtrl.Finish()
+		})
+	}
+}
+
 func TestAttachDisk(t *testing.T) {
 	testCases := []struct {
 		name     string