From e350739ecb0f3d0916005e108e99f0c5219f1316 Mon Sep 17 00:00:00 2001
From: DanHam <DanHam@users.noreply.github.com>
Date: Mon, 23 Jan 2017 22:25:39 +0000
Subject: [PATCH] post-processor/shell-local: expand tests to align with shell
 provisioner

---
 .../shell-local/post-processor_test.go        | 53 +++++++++++++++----
 provisioner/shell/provisioner_test.go         |  1 -
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/post-processor/shell-local/post-processor_test.go b/post-processor/shell-local/post-processor_test.go
index 3532b9cc6..c5681694a 100644
--- a/post-processor/shell-local/post-processor_test.go
+++ b/post-processor/shell-local/post-processor_test.go
@@ -198,22 +198,57 @@ func TestPostProcessorPrepare_EnvironmentVars(t *testing.T) {
 	if err != nil {
 		t.Fatalf("should not have error: %s", err)
 	}
+
+	// Test when the env variable value contains an equals sign
+	config["environment_vars"] = []string{"good=withequals=true"}
+	p = new(PostProcessor)
+	err = p.Configure(config)
+	if err != nil {
+		t.Fatalf("should not have error: %s", err)
+	}
+
+	// Test when the env variable value starts with an equals sign
+	config["environment_vars"] = []string{"good==true"}
+	p = new(PostProcessor)
+	err = p.Configure(config)
+	if err != nil {
+		t.Fatalf("should not have error: %s", err)
+	}
 }
 
-func TestPostProcessorQuote_EnvironmentVars(t *testing.T) {
+func TestPostProcessor_createFlattenedEnvVars(t *testing.T) {
+	var flattenedEnvVars string
 	config := testConfig()
 
-	config["environment_vars"] = []string{"keyone=valueone", "keytwo=value\ntwo"}
+	userEnvVarTests := [][]string{
+		{},                     // No user env var
+		{"FOO=bar"},            // Single user env var
+		{"FOO=bar's"},          // User env var with single quote in value
+		{"FOO=bar", "BAZ=qux"}, // Multiple user env vars
+		{"FOO=bar=baz"},        // User env var with value containing equals
+		{"FOO==bar"},           // User env var with value starting with equals
+	}
+	expected := []string{
+		`PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+		`FOO='bar' PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+		`FOO='bar'"'"'s' PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+		`BAZ='qux' FOO='bar' PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+		`FOO='bar=baz' PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+		`FOO='=bar' PACKER_BUILDER_TYPE='iso' PACKER_BUILD_NAME='vmware' `,
+	}
+
 	p := new(PostProcessor)
 	p.Configure(config)
 
-	expectedValue := "keyone='valueone'"
-	if p.config.Vars[0] != expectedValue {
-		t.Fatalf("%s should be equal to %s", p.config.Vars[0], expectedValue)
-	}
+	// Defaults provided by Packer
+	p.config.PackerBuildName = "vmware"
+	p.config.PackerBuilderType = "iso"
 
-	expectedValue = "keytwo='value\ntwo'"
-	if p.config.Vars[1] != expectedValue {
-		t.Fatalf("%s should be equal to %s", p.config.Vars[1], expectedValue)
+	for i, expectedValue := range expected {
+		p.config.Vars = userEnvVarTests[i]
+		flattenedEnvVars = p.createFlattenedEnvVars()
+		if flattenedEnvVars != expectedValue {
+			t.Fatalf("expected flattened env vars to be: %s, got %s.", expectedValue, flattenedEnvVars)
+		}
 	}
 }
diff --git a/provisioner/shell/provisioner_test.go b/provisioner/shell/provisioner_test.go
index 007e5e7c3..c2bcb455d 100644
--- a/provisioner/shell/provisioner_test.go
+++ b/provisioner/shell/provisioner_test.go
@@ -247,7 +247,6 @@ func TestProvisionerPrepare_EnvironmentVars(t *testing.T) {
 	if err != nil {
 		t.Fatalf("should not have error: %s", err)
 	}
-
 }
 
 func TestProvisioner_createFlattenedEnvVars(t *testing.T) {