From 3edf49ef039b23cc2ee0edd39679058fdebb8a51 Mon Sep 17 00:00:00 2001
From: Justin Barrick <justin.m.barrick@gmail.com>
Date: Thu, 28 Feb 2019 23:49:59 -0800
Subject: [PATCH 1/2] Fix support for volume snapshots by setting snapshot id
 on volume creation.

---
 driver/controller.go                |  9 +++++++
 test/kubernetes/integration_test.go | 38 ++++++++++++++++++++++++++---
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/driver/controller.go b/driver/controller.go
index 5cfd1e70b..aa93836a2 100644
--- a/driver/controller.go
+++ b/driver/controller.go
@@ -157,6 +157,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.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..dd5c8bb39 100644
--- a/test/kubernetes/integration_test.go
+++ b/test/kubernetes/integration_test.go
@@ -358,6 +358,22 @@ func TestSnapshot_Create(t *testing.T) {
 			Name: "my-csi-app-2",
 		},
 		Spec: v1.PodSpec{
+			InitContainers: []v1.Container{
+				{
+					Name:  "my-csi",
+					Image: "busybox",
+					VolumeMounts: []v1.VolumeMount{
+						{
+							MountPath: "/data",
+							Name:      volumeName,
+						},
+					},
+					Command: []string{
+						"sh", "-c",
+						"echo testcanary > /data/canary",
+					},
+				},
+			},
 			Containers: []v1.Container{
 				{
 					Name:  "my-csi-app",
@@ -369,8 +385,7 @@ func TestSnapshot_Create(t *testing.T) {
 						},
 					},
 					Command: []string{
-						"sleep",
-						"1000000",
+						"sleep", "1000000",
 					},
 				},
 			},
@@ -476,6 +491,22 @@ func TestSnapshot_Create(t *testing.T) {
 			Name: "my-csi-app-2-restored",
 		},
 		Spec: v1.PodSpec{
+			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 +518,7 @@ func TestSnapshot_Create(t *testing.T) {
 						},
 					},
 					Command: []string{
-						"sleep",
-						"1000000",
+						"sleep", "1000000",
 					},
 				},
 			},

From 71383463d1b8bfdf08f48b2fbef0d80ea9a821e0 Mon Sep 17 00:00:00 2001
From: Justin Barrick <justin.m.barrick@gmail.com>
Date: Mon, 4 Mar 2019 19:37:46 -0800
Subject: [PATCH 2/2] Address PR comments

---
 driver/controller.go                |  4 ++--
 test/kubernetes/integration_test.go | 24 +++++++-----------------
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/driver/controller.go b/driver/controller.go
index aa93836a2..a5cff625e 100644
--- a/driver/controller.go
+++ b/driver/controller.go
@@ -161,8 +161,8 @@ func (d *Driver) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest)
 	if contentSource != nil {
 		snapshot := contentSource.GetSnapshot()
 		if snapshot != nil {
-			ll.Info("using snapshot as volume source")
-			volumeReq.SnapshotID = snapshot.GetId()
+			ll.WithField("snapshot_id", snapshot.GetSnapshotId()).Info("using snapshot as volume source")
+			volumeReq.SnapshotID = snapshot.GetSnapshotId()
 		}
 	}
 
diff --git a/test/kubernetes/integration_test.go b/test/kubernetes/integration_test.go
index dd5c8bb39..ea50f420a 100644
--- a/test/kubernetes/integration_test.go
+++ b/test/kubernetes/integration_test.go
@@ -358,22 +358,6 @@ func TestSnapshot_Create(t *testing.T) {
 			Name: "my-csi-app-2",
 		},
 		Spec: v1.PodSpec{
-			InitContainers: []v1.Container{
-				{
-					Name:  "my-csi",
-					Image: "busybox",
-					VolumeMounts: []v1.VolumeMount{
-						{
-							MountPath: "/data",
-							Name:      volumeName,
-						},
-					},
-					Command: []string{
-						"sh", "-c",
-						"echo testcanary > /data/canary",
-					},
-				},
-			},
 			Containers: []v1.Container{
 				{
 					Name:  "my-csi-app",
@@ -385,7 +369,8 @@ func TestSnapshot_Create(t *testing.T) {
 						},
 					},
 					Command: []string{
-						"sleep", "1000000",
+						"sh", "-c",
+						"echo testcanary > /data/canary && sleep 1000000",
 					},
 				},
 			},
@@ -491,6 +476,11 @@ 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",