From de4839b66ead6fbe3de8db538cdb4c87f18a1e4c Mon Sep 17 00:00:00 2001
From: Paul Meyer <paulmey@microsoft.com>
Date: Wed, 29 Apr 2020 19:14:15 +0000
Subject: [PATCH] Use all snapshots when creating a shared image

---
 .../step_create_shared_image_version.go       | 26 ++++++++++++++++---
 .../step_create_shared_image_version_test.go  | 26 ++++++++++++++++---
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/builder/azure/chroot/step_create_shared_image_version.go b/builder/azure/chroot/step_create_shared_image_version.go
index 83543f1e6..df6afdd6c 100644
--- a/builder/azure/chroot/step_create_shared_image_version.go
+++ b/builder/azure/chroot/step_create_shared_image_version.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"sort"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-12-01/compute"
@@ -22,11 +23,11 @@ type StepCreateSharedImageVersion struct {
 func (s *StepCreateSharedImageVersion) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
 	azcli := state.Get("azureclient").(client.AzureClientSet)
 	ui := state.Get("ui").(packer.Ui)
-	osDiskSnapshotResourceID := state.Get(stateBagKey_Snapshotset).(Diskset)[-1].String()
+	snapshotset := state.Get(stateBagKey_Snapshotset).(Diskset)
 
-	ui.Say(fmt.Sprintf("Creating image version %s\n   using %s for os disk.",
+	ui.Say(fmt.Sprintf("Creating image version %s\n   using %q for os disk.",
 		s.Destination.ResourceID(azcli.SubscriptionID()),
-		osDiskSnapshotResourceID))
+		snapshotset.OS()))
 
 	var targetRegions []compute.TargetRegion
 	// transform target regions to API objects
@@ -44,7 +45,7 @@ func (s *StepCreateSharedImageVersion) Run(ctx context.Context, state multistep.
 		GalleryImageVersionProperties: &compute.GalleryImageVersionProperties{
 			StorageProfile: &compute.GalleryImageVersionStorageProfile{
 				OsDiskImage: &compute.GalleryOSDiskImage{
-					Source:      &compute.GalleryArtifactVersionSource{ID: &osDiskSnapshotResourceID},
+					Source:      &compute.GalleryArtifactVersionSource{ID: to.StringPtr(snapshotset.OS().String())},
 					HostCaching: compute.HostCaching(s.OSDiskCacheType),
 				},
 			},
@@ -55,6 +56,23 @@ func (s *StepCreateSharedImageVersion) Run(ctx context.Context, state multistep.
 		},
 	}
 
+	var datadisks []compute.GalleryDataDiskImage
+	for lun, resource := range snapshotset {
+		if lun != -1 {
+			datadisks = append(datadisks, compute.GalleryDataDiskImage{
+				Lun:    to.Int32Ptr(lun),
+				Source: &compute.GalleryArtifactVersionSource{ID: to.StringPtr(resource.String())},
+			})
+		}
+	}
+	if datadisks != nil {
+		// sort by lun
+		sort.Slice(datadisks, func(i, j int) bool {
+			return *datadisks[i].Lun < *datadisks[j].Lun
+		})
+		imageVersion.GalleryImageVersionProperties.StorageProfile.DataDiskImages = &datadisks
+	}
+
 	f, err := azcli.GalleryImageVersionsClient().CreateOrUpdate(
 		ctx,
 		s.Destination.ResourceGroup,
diff --git a/builder/azure/chroot/step_create_shared_image_version_test.go b/builder/azure/chroot/step_create_shared_image_version_test.go
index 09a750c48..c26411f3b 100644
--- a/builder/azure/chroot/step_create_shared_image_version_test.go
+++ b/builder/azure/chroot/step_create_shared_image_version_test.go
@@ -25,6 +25,7 @@ func TestStepCreateSharedImageVersion_Run(t *testing.T) {
 	tests := []struct {
 		name            string
 		fields          fields
+		snapshotset     Diskset
 		want            multistep.StepAction
 		expectedPutBody string
 	}{
@@ -47,6 +48,11 @@ func TestStepCreateSharedImageVersion_Run(t *testing.T) {
 				},
 				Location: "region2",
 			},
+			snapshotset: diskset(
+				"/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/osdisksnapshot",
+				"/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot0",
+				"/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot1",
+				"/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot2"),
 			expectedPutBody: `{
 				"location": "region2",
 				"properties": {
@@ -63,9 +69,23 @@ func TestStepCreateSharedImageVersion_Run(t *testing.T) {
 					"storageProfile": {
 						"osDiskImage": {
 							"source": {
-								"id": "/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/snapshot1"
+								"id": "/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/osdisksnapshot"
 							}
-						}
+						},
+						"dataDiskImages": [
+							{
+								"lun": 0,
+								"source": { "id": "/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot0" }
+							},
+							{
+								"lun": 1,
+								"source": { "id": "/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot1" }
+							},
+							{
+								"lun": 2,
+								"source": { "id": "/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/datadisksnapshot2" }
+							}
+						]
 					}
 				}
 			}`,
@@ -94,7 +114,7 @@ func TestStepCreateSharedImageVersion_Run(t *testing.T) {
 			GalleryImageVersionsClientMock: m,
 		})
 		state.Put("ui", packer.TestUi(t))
-		state.Put(stateBagKey_Snapshotset, diskset("/subscriptions/12345/resourceGroups/group1/providers/Microsoft.Compute/snapshots/snapshot1"))
+		state.Put(stateBagKey_Snapshotset, tt.snapshotset)
 
 		t.Run(tt.name, func(t *testing.T) {
 			s := &StepCreateSharedImageVersion{