From c167ef91ced84ad1de4f17bc1cefd349414ca102 Mon Sep 17 00:00:00 2001
From: Paul Meyer <paulmey@microsoft.com>
Date: Wed, 25 Mar 2020 18:24:18 +0000
Subject: [PATCH] Move PIR :latest resolution to step

---
 builder/azure/chroot/builder.go               | 12 ++--
 .../step_resolve_plaform_image_version.go     | 45 ++++++++++++++
 ...step_resolve_plaform_image_version_test.go | 61 +++++++++++++++++++
 3 files changed, 110 insertions(+), 8 deletions(-)
 create mode 100644 builder/azure/chroot/step_resolve_plaform_image_version.go
 create mode 100644 builder/azure/chroot/step_resolve_plaform_image_version_test.go

diff --git a/builder/azure/chroot/builder.go b/builder/azure/chroot/builder.go
index 7d95c2097..4995a860e 100644
--- a/builder/azure/chroot/builder.go
+++ b/builder/azure/chroot/builder.go
@@ -27,7 +27,6 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
 	"github.com/Azure/go-autorest/autorest/azure"
-	"github.com/Azure/go-autorest/autorest/to"
 )
 
 // BuilderId is the unique ID for this builder
@@ -356,13 +355,10 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
 
 			if pi, err := client.ParsePlatformImageURN(b.config.Source); err == nil {
 				if strings.EqualFold(pi.Version, "latest") {
-
-					vmi, err := azcli.VirtualMachineImagesClient().GetLatest(ctx, pi.Publisher, pi.Offer, pi.Sku, info.Location)
-					if err != nil {
-						return nil, fmt.Errorf("error retieving latest version of %q: %v", b.config.Source, err)
-					}
-					pi.Version = to.String(vmi.Name)
-					log.Println("Resolved latest version of source image:", pi.Version)
+					steps = append(steps, &StepResolvePlatformImageVersion{
+						PlatformImage: pi,
+						Location:      info.Location,
+					})
 				}
 				steps = append(steps,
 
diff --git a/builder/azure/chroot/step_resolve_plaform_image_version.go b/builder/azure/chroot/step_resolve_plaform_image_version.go
new file mode 100644
index 000000000..77a2e7531
--- /dev/null
+++ b/builder/azure/chroot/step_resolve_plaform_image_version.go
@@ -0,0 +1,45 @@
+package chroot
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strings"
+
+	"github.com/Azure/go-autorest/autorest/to"
+	"github.com/hashicorp/packer/builder/azure/common/client"
+	"github.com/hashicorp/packer/helper/multistep"
+	"github.com/hashicorp/packer/packer"
+)
+
+// StepResolvePlatformImageVersion resolves the exact PIR version when the version is 'latest'
+type StepResolvePlatformImageVersion struct {
+	*client.PlatformImage
+	Location string
+}
+
+// Run retrieves all available versions of a PIR image and stores the latest in the PlatformImage
+func (pi *StepResolvePlatformImageVersion) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
+	ui := state.Get("ui").(packer.Ui)
+
+	if strings.EqualFold(pi.Version, "latest") {
+		azcli := state.Get("azureclient").(client.AzureClientSet)
+
+		vmi, err := azcli.VirtualMachineImagesClient().GetLatest(ctx, pi.Publisher, pi.Offer, pi.Sku, pi.Location)
+		if err != nil {
+			log.Printf("StepResolvePlatformImageVersion.Run: error: %+v", err)
+			err := fmt.Errorf("error retieving latest version of %q: %v", pi.URN(), err)
+			state.Put("error", err)
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+		pi.Version = to.String(vmi.Name)
+		ui.Say("Resolved latest version of source image: " + pi.Version)
+	} else {
+		ui.Say("Nothing to do, version is not 'latest'")
+	}
+
+	return multistep.ActionContinue
+}
+
+func (*StepResolvePlatformImageVersion) Cleanup(multistep.StateBag) {}
diff --git a/builder/azure/chroot/step_resolve_plaform_image_version_test.go b/builder/azure/chroot/step_resolve_plaform_image_version_test.go
new file mode 100644
index 000000000..8adebf2ee
--- /dev/null
+++ b/builder/azure/chroot/step_resolve_plaform_image_version_test.go
@@ -0,0 +1,61 @@
+package chroot
+
+import (
+	"context"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-03-01/compute"
+	"github.com/Azure/go-autorest/autorest"
+	"github.com/hashicorp/packer/builder/azure/common/client"
+	"github.com/hashicorp/packer/helper/multistep"
+)
+
+func TestStepResolvePlatformImageVersion_Run(t *testing.T) {
+
+	pi := &StepResolvePlatformImageVersion{
+		PlatformImage: &client.PlatformImage{
+			Version: "latest",
+		}}
+
+	m := compute.NewVirtualMachineImagesClient("subscriptionId")
+	m.Sender = autorest.SenderFunc(func(r *http.Request) (*http.Response, error) {
+		if !strings.Contains(r.URL.String(), "%24orderby=name+desc") {
+			t.Errorf("Expected url to use odata based sorting, but got %q", r.URL.String())
+		}
+		return &http.Response{
+			Request: r,
+			Body: ioutil.NopCloser(strings.NewReader(
+				`[
+					{"name":"1.2.3"},
+					{"name":"4.5.6"}
+				]`)),
+			StatusCode: 200,
+		}, nil
+	})
+
+	state := new(multistep.BasicStateBag)
+	state.Put("azureclient", &client.AzureClientSetMock{
+		VirtualMachineImagesClientMock: client.VirtualMachineImagesClient{m},
+	})
+
+	ui, getErrs := testUI()
+	state.Put("ui", ui)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	got := pi.Run(ctx, state)
+	if got != multistep.ActionContinue {
+		t.Errorf("Expected 'continue', but got %q", got)
+	}
+
+	if pi.PlatformImage.Version != "1.2.3" {
+		t.Errorf("Expected version '1.2.3', but got %q", pi.PlatformImage.Version)
+	}
+
+	_ = getErrs
+}