From 5caf374a4ce57de7a9b332846aeee19c104880ed Mon Sep 17 00:00:00 2001
From: Fatih Arslan <ftharsln@gmail.com>
Date: Wed, 20 Mar 2019 13:45:11 +0300
Subject: [PATCH] Merge pull request #129 from justinbarrick/release-0.4.0

Fix support for volume snapshots by setting snapshot id on volume creation.
---
 driver/controller.go                |  9 +++++++++
 test/kubernetes/integration_test.go | 28 ++++++++++++++++++++++++----
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/driver/controller.go b/driver/controller.go
index 113ad2e10..463188e90 100644
--- a/driver/controller.go
+++ b/driver/controller.go
@@ -156,6 +156,15 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 		volumeReq.Tags = append(volumeReq.Tags, d.doTag)
 	}
 
+	contentSource := req.GetVolumeContentSource()
+	if contentSource != nil {
+		snapshot := contentSource.GetSnapshot()
+		if snapshot != nil {
+			ll.WithField("snapshot_id", snapshot.GetId()).Info("using snapshot as volume source")
+			volumeReq.SnapshotID = snapshot.GetId()
+		}
+	}
+
 	ll.Info("checking volume limit")
 	if err := d.checkLimit(ctx); err != nil {
 		return nil, err
diff --git a/test/kubernetes/integration_test.go b/test/kubernetes/integration_test.go
index e63ec2d66..ea50f420a 100644
--- a/test/kubernetes/integration_test.go
+++ b/test/kubernetes/integration_test.go
@@ -369,8 +369,8 @@ func TestSnapshot_Create(t *testing.T) {
 						},
 					},
 					Command: []string{
-						"sleep",
-						"1000000",
+						"sh", "-c",
+						"echo testcanary > /data/canary && sleep 1000000",
 					},
 				},
 			},
@@ -476,6 +476,27 @@ func TestSnapshot_Create(t *testing.T) {
 			Name: "my-csi-app-2-restored",
 		},
 		Spec: v1.PodSpec{
+			// This init container verifies that the /data/canary file is present.
+			// If it is not, then the volume was not properly restored.
+			// waitForPod only waits for the pod to enter the running state, so will not
+			// detect any failures after that, so this has to be an InitContainer so that
+			// the pod never enters the running state if it fails.
+			InitContainers: []v1.Container{
+				{
+					Name:  "my-csi",
+					Image: "busybox",
+					VolumeMounts: []v1.VolumeMount{
+						{
+							MountPath: "/data",
+							Name:      volumeName,
+						},
+					},
+					Command: []string{
+						"cat",
+						"/data/canary",
+					},
+				},
+			},
 			Containers: []v1.Container{
 				{
 					Name:  "my-csi-app",
@@ -487,8 +508,7 @@ func TestSnapshot_Create(t *testing.T) {
 						},
 					},
 					Command: []string{
-						"sleep",
-						"1000000",
+						"sleep", "1000000",
 					},
 				},
 			},