From b2d1675d399d2a1f0c177b12bd240019ee64895a Mon Sep 17 00:00:00 2001
From: Amrita Dutta <amdut@microsoft.com>
Date: Wed, 7 Nov 2018 03:23:17 +0000
Subject: [PATCH] Added tests

---
 .../arm/step_snapshot_data_disks_test.go      | 69 ++++++++++++++++++
 .../azure/arm/step_snapshot_os_disk_test.go   | 70 +++++++++++++++++++
 2 files changed, 139 insertions(+)
 create mode 100644 builder/azure/arm/step_snapshot_data_disks_test.go
 create mode 100644 builder/azure/arm/step_snapshot_os_disk_test.go

diff --git a/builder/azure/arm/step_snapshot_data_disks_test.go b/builder/azure/arm/step_snapshot_data_disks_test.go
new file mode 100644
index 000000000..0eefa924d
--- /dev/null
+++ b/builder/azure/arm/step_snapshot_data_disks_test.go
@@ -0,0 +1,69 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"testing"
+)
+
+func TestStepSnapshotDataDisksShouldFailIfSnapshotFails(t *testing.T) {
+	var testSubject = &StepSnapshotDataDisks{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return fmt.Errorf("!! Unit Test FAIL !!")
+		},
+		say:   func(message string) {},
+		error: func(e error) {},
+	}
+
+	stateBag := createTestStateBagStepSnapshotDataDisks()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
+	}
+}
+
+func TestStepSnapshotDataDisksShouldPassIfSnapshotPasses(t *testing.T) {
+	var testSubject = &StepSnapshotDataDisks{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return nil
+		},
+		say:   func(message string) {},
+		error: func(e error) {},
+	}
+
+	stateBag := createTestStateBagStepSnapshotDataDisks()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionContinue {
+		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == true {
+		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
+	}
+}
+
+func createTestStateBagStepSnapshotDataDisks() multistep.StateBag {
+	stateBag := new(multistep.BasicStateBag)
+
+	stateBag.Put(constants.ArmResourceGroupName, "Unit Test: ResourceGroupName")
+	stateBag.Put(constants.ArmLocation, "Unit Test: Location")
+
+	value := "Unit Test: Tags"
+	tags := map[string]*string{
+		"tag01": &value,
+	}
+	stateBag.Put(constants.ArmTags, tags)
+
+	stateBag.Put(constants.ArmAdditionalDiskVhds, []string{"subscriptions/123-456-789/resourceGroups/existingresourcegroup/providers/Microsoft.Compute/disks/osdisk"})
+	stateBag.Put(constants.ArmManagedImageDataDiskSnapshotPrefix, "Unit Test: ManagedImageDataDiskSnapshotPrefix")
+
+	return stateBag
+}
diff --git a/builder/azure/arm/step_snapshot_os_disk_test.go b/builder/azure/arm/step_snapshot_os_disk_test.go
new file mode 100644
index 000000000..08f5038ac
--- /dev/null
+++ b/builder/azure/arm/step_snapshot_os_disk_test.go
@@ -0,0 +1,70 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"testing"
+)
+
+func TestStepSnapshotOSDiskShouldFailIfSnapshotFails(t *testing.T) {
+	var testSubject = &StepSnapshotOSDisk{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return fmt.Errorf("!! Unit Test FAIL !!")
+		},
+		say:   func(message string) {},
+		error: func(e error) {},
+	}
+
+	stateBag := createTestStateBagStepSnapshotOSDisk()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
+	}
+}
+
+func TestStepSnapshotOSDiskShouldPassIfSnapshotPasses(t *testing.T) {
+	var testSubject = &StepSnapshotOSDisk{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return nil
+		},
+		say:   func(message string) {},
+		error: func(e error) {},
+	}
+
+	stateBag := createTestStateBagStepSnapshotOSDisk()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionContinue {
+		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == true {
+		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
+	}
+}
+
+func createTestStateBagStepSnapshotOSDisk() multistep.StateBag {
+	stateBag := new(multistep.BasicStateBag)
+
+	stateBag.Put(constants.ArmResourceGroupName, "Unit Test: ResourceGroupName")
+	stateBag.Put(constants.ArmLocation, "Unit Test: Location")
+
+	value := "Unit Test: Tags"
+	tags := map[string]*string{
+		"tag01": &value,
+	}
+
+	stateBag.Put(constants.ArmTags, tags)
+
+	stateBag.Put(constants.ArmOSDiskVhd, "subscriptions/123-456-789/resourceGroups/existingresourcegroup/providers/Microsoft.Compute/disks/osdisk")
+	stateBag.Put(constants.ArmManagedImageOSDiskSnapshotName, "Unit Test: ManagedImageOSDiskSnapshotName")
+
+	return stateBag
+}
\ No newline at end of file