From f1e966401272b0594695d9e7bd3d5978dcfcfefc Mon Sep 17 00:00:00 2001
From: krisko <tomasz.krzyszkowski@tradeshift.com>
Date: Tue, 23 Apr 2019 20:49:41 +0200
Subject: [PATCH] Googlebuilder: Add image key encryption

---
 builder/googlecompute/config_test.go          | 15 +++++++++++
 builder/googlecompute/driver_mock.go          | 27 ++++++++++++-------
 .../googlecompute/step_create_image_test.go   |  1 +
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/builder/googlecompute/config_test.go b/builder/googlecompute/config_test.go
index 77086a413..755050b93 100644
--- a/builder/googlecompute/config_test.go
+++ b/builder/googlecompute/config_test.go
@@ -156,6 +156,21 @@ func TestConfigPrepare(t *testing.T) {
 			"foo bar",
 			true,
 		},
+		{
+			"image_encryption_key",
+			map[string]string{"kmsKeyName": "foo"},
+			false,
+		},
+		{
+			"image_encryption_key",
+			map[string]string{"No such key": "foo"},
+			true,
+		},
+		{
+			"image_encryption_key",
+			map[string]string{"kmsKeyName": "foo", "RawKey": "foo"},
+			false,
+		},
 		{
 			"scopes",
 			[]string{},
diff --git a/builder/googlecompute/driver_mock.go b/builder/googlecompute/driver_mock.go
index 720f9d60b..103a7ed47 100644
--- a/builder/googlecompute/driver_mock.go
+++ b/builder/googlecompute/driver_mock.go
@@ -1,6 +1,10 @@
 package googlecompute
 
-import "fmt"
+import (
+	"fmt"
+
+	compute "google.golang.org/api/compute/v1"
+)
 
 // DriverMock is a Driver implementation that is a mocked out so that
 // it can be used for tests.
@@ -9,6 +13,7 @@ type DriverMock struct {
 	CreateImageDesc            string
 	CreateImageFamily          string
 	CreateImageLabels          map[string]string
+	CreateImageEncryptionKey   *compute.CustomerEncryptionKey
 	CreateImageLicenses        []string
 	CreateImageZone            string
 	CreateImageDisk            string
@@ -82,14 +87,16 @@ type DriverMock struct {
 	WaitForInstanceErrCh <-chan error
 }
 
-func (d *DriverMock) CreateImage(name, description, family, zone, disk string, image_labels map[string]string, image_licenses []string) (<-chan *Image, <-chan error) {
-	d.CreateImageName = name
-	d.CreateImageDesc = description
-	d.CreateImageFamily = family
-	d.CreateImageLabels = image_labels
-	d.CreateImageLicenses = image_licenses
-	d.CreateImageZone = zone
-	d.CreateImageDisk = disk
+func (d *DriverMock) CreateImage(config Config) (<-chan *Image, <-chan error) {
+	d.CreateImageName = config.GetImageName()
+	d.CreateImageDesc = config.GetImageDescription()
+	d.CreateImageFamily = config.GetImageFamily()
+	d.CreateImageLabels = config.GetImageLabels()
+	d.CreateImageLicenses = config.GetImageLicenses()
+	d.CreateImageZone = config.GetZone()
+	d.CreateImageDisk = config.GetDiskName()
+	d.CreateImageEncryptionKey = config.GetImageEncryptionKey()
+
 	if d.CreateImageResultProjectId == "" {
 		d.CreateImageResultProjectId = "test"
 	}
@@ -108,7 +115,7 @@ func (d *DriverMock) CreateImage(name, description, family, zone, disk string, i
 		ch <- &Image{
 			Labels:    d.CreateImageLabels,
 			Licenses:  d.CreateImageLicenses,
-			Name:      name,
+			Name:      d.CreateImageName,
 			ProjectId: d.CreateImageResultProjectId,
 			SelfLink:  d.CreateImageResultSelfLink,
 			SizeGb:    d.CreateImageResultSizeGb,
diff --git a/builder/googlecompute/step_create_image_test.go b/builder/googlecompute/step_create_image_test.go
index e870eba42..0f1458124 100644
--- a/builder/googlecompute/step_create_image_test.go
+++ b/builder/googlecompute/step_create_image_test.go
@@ -47,6 +47,7 @@ func TestStepCreateImage(t *testing.T) {
 	assert.Equal(t, d.CreateImageDisk, c.DiskName, "Incorrect disk passed to driver.")
 	assert.Equal(t, d.CreateImageLabels, c.ImageLabels, "Incorrect image_labels passed to driver.")
 	assert.Equal(t, d.CreateImageLicenses, c.ImageLicenses, "Incorrect image_licenses passed to driver.")
+	assert.Equal(t, d.CreateImageEncryptionKey, c.ImageEncryptionKey, "Incorrect image_encryption_key passed to driver.")
 }
 
 func TestStepCreateImage_errorOnChannel(t *testing.T) {