From 7416d2f26fec8cf7cd4f8e989d8e11b1bef9a6a2 Mon Sep 17 00:00:00 2001
From: Cheng Pan <chengpan@amazon.com>
Date: Wed, 24 Oct 2018 15:22:39 -0700
Subject: [PATCH] Add unit test

---
 pkg/cloud/cloud.go      | 68 ++++++++++++++++++++++++++++++++++++++---
 pkg/cloud/cloud_test.go |  5 +++
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/pkg/cloud/cloud.go b/pkg/cloud/cloud.go
index 47f69c2324..783ba2ff29 100644
--- a/pkg/cloud/cloud.go
+++ b/pkg/cloud/cloud.go
@@ -244,16 +244,23 @@ func (c *cloud) CreateDisk(ctx context.Context, volumeName string, diskOptions *
 		return nil, fmt.Errorf("disk size was not returned by CreateVolume")
 	}
 
+	if len(diskOptions.KmsKeyID) > 0 {
+		err := c.waitForCreate(ctx, volumeID)
+		if err != nil {
+			if isAWSErrorVolumeNotFound(err) {
+				return nil, fmt.Errorf("failed to create encrypted volume: the volume disappeared after creation, most likely due to inaccessible KMS encryption key")
+			}
+		}
+	}
+
 	return &Disk{CapacityGiB: size, VolumeID: volumeID, AvailabilityZone: zone}, nil
 }
 
 func (c *cloud) DeleteDisk(ctx context.Context, volumeID string) (bool, error) {
 	request := &ec2.DeleteVolumeInput{VolumeId: &volumeID}
 	if _, err := c.ec2.DeleteVolumeWithContext(ctx, request); err != nil {
-		if awsErr, ok := err.(awserr.Error); ok {
-			if awsErr.Code() == "InvalidVolume.NotFound" {
-				return false, ErrNotFound
-			}
+		if isAWSErrorVolumeNotFound(err) {
+			return false, ErrNotFound
 		}
 		return false, fmt.Errorf("DeleteDisk could not delete volume: %v", err)
 	}
@@ -496,3 +503,56 @@ func (c *cloud) waitForAttachmentState(ctx context.Context, volumeID, state stri
 
 	return wait.ExponentialBackoff(backoff, verifyVolumeFunc)
 }
+
+// waitForCreate waits for volume to be created for encrypted volume only
+// it polls for created volume to check it has not been silently removed by AWS.
+// On a random AWS account (shared among several developers) it took 4s on average.
+func (c *cloud) waitForCreate(ctx context.Context, volumeID string) error {
+	var (
+		checkInterval = 1 * time.Second
+		checkTimeout  = 30 * time.Second
+	)
+
+	request := &ec2.DescribeVolumesInput{
+		VolumeIds: []*string{
+			aws.String(volumeID),
+		},
+	}
+
+	err := wait.Poll(checkInterval, checkTimeout, func() (done bool, err error) {
+		vol, err := c.getVolume(ctx, request)
+		if err != nil {
+			return true, err
+		}
+		if vol.State != nil {
+			switch *vol.State {
+			case "available":
+				// The volume is Available, it won't be deleted now.
+				return true, nil
+			case "creating":
+				return false, nil
+			default:
+				return true, fmt.Errorf("unexpected State of newly created AWS EBS volume %s: %q", volumeID, *vol.State)
+			}
+		}
+		return false, nil
+	})
+
+	return err
+}
+
+// Helper function for describeVolume callers. Tries to retype given error to AWS error
+// and returns true in case the AWS error is "InvalidVolume.NotFound", false otherwise
+func isAWSErrorVolumeNotFound(err error) bool {
+	if err == nil {
+		return false
+	}
+	if awsError, ok := err.(awserr.Error); ok {
+		// https://docs.aws.amazon.com/AWSEC2/latest/APIReference/errors-overview.html
+		if awsError.Code() == "InvalidVolume.NotFound" {
+			return true
+		}
+	}
+
+	return false
+}
diff --git a/pkg/cloud/cloud_test.go b/pkg/cloud/cloud_test.go
index f475be5eef..f9179d80d1 100644
--- a/pkg/cloud/cloud_test.go
+++ b/pkg/cloud/cloud_test.go
@@ -106,12 +106,17 @@ func TestCreateDisk(t *testing.T) {
 			vol = &ec2.Volume{
 				VolumeId: aws.String(tc.diskOptions.Tags[VolumeNameTagKey]),
 				Size:     aws.Int64(util.BytesToGiB(tc.diskOptions.CapacityBytes)),
+				State:    aws.String("available"),
 			}
 		}
 
 		ctx := context.Background()
 		mockEC2.EXPECT().CreateVolumeWithContext(gomock.Eq(ctx), gomock.Any()).Return(vol, tc.expErr)
 
+		if tc.diskOptions.Encrypted {
+			mockEC2.EXPECT().DescribeVolumesWithContext(gomock.Eq(ctx), gomock.Any()).Return(&ec2.DescribeVolumesOutput{Volumes: []*ec2.Volume{vol}}, nil)
+		}
+
 		disk, err := c.CreateDisk(ctx, tc.volumeName, tc.diskOptions)
 		if err != nil {
 			if tc.expErr == nil {