From 37bd1cd3088d7fdfb47b610e1eb2abf430890606 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 15 Apr 2013 20:26:38 -0700
Subject: [PATCH] Test setting the UI for tests

---
 packer.go                  |  2 +-
 packer/environment.go      | 37 +++++++++++++++++++++++++++++++++----
 packer/environment_test.go | 32 ++++++++++++++++++++++++++------
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/packer.go b/packer.go
index 67259a9cc..015bd4ff1 100644
--- a/packer.go
+++ b/packer.go
@@ -7,6 +7,6 @@ import (
 )
 
 func main() {
-	env := packer.NewEnvironment()
+	env := packer.NewEnvironment(nil)
 	os.Exit(env.Cli(os.Args[1:]))
 }
diff --git a/packer/environment.go b/packer/environment.go
index 083f7a418..d38a78fd3 100644
--- a/packer/environment.go
+++ b/packer/environment.go
@@ -39,14 +39,43 @@ type Environment struct {
 	ui      Ui
 }
 
+// This struct configures new environments.
+type EnvironmentConfig struct {
+	builder map[string]Builder
+	command map[string]Command
+	ui      Ui
+}
+
 // This creates a new environment
-func NewEnvironment() *Environment {
+func NewEnvironment(config *EnvironmentConfig) *Environment {
 	env := &Environment{}
 	env.builder = make(map[string]Builder)
 	env.command = make(map[string]Command)
-	env.command["build"] = new(buildCommand)
-	env.command["version"] = new(versionCommand)
-	env.ui = &ReaderWriterUi{os.Stdin, os.Stdout}
+
+	if config != nil {
+		for k, v := range config.builder {
+			env.builder[k] = v
+		}
+
+		for k, v := range config.command {
+			env.command[k] = v
+		}
+
+		env.ui = config.ui
+	}
+
+	if _, ok := env.command["build"]; !ok {
+		env.command["build"] = new(buildCommand)
+	}
+
+	if _, ok := env.command["version"]; !ok {
+		env.command["version"] = new(versionCommand)
+	}
+
+	if env.ui == nil {
+		env.ui = &ReaderWriterUi{os.Stdin, os.Stdout}
+	}
+
 	return env
 }
 
diff --git a/packer/environment_test.go b/packer/environment_test.go
index fc62c692c..281a18a05 100644
--- a/packer/environment_test.go
+++ b/packer/environment_test.go
@@ -1,13 +1,20 @@
 package packer
 
 import (
+	"bytes"
 	"cgl.tideland.biz/asserts"
 	"os"
 	"testing"
 )
 
 func testEnvironment() *Environment {
-	return NewEnvironment()
+	config := &EnvironmentConfig{}
+	config.ui = &ReaderWriterUi{
+		new(bytes.Buffer),
+		new(bytes.Buffer),
+	}
+
+	return NewEnvironment(config)
 }
 
 func TestEnvironment_Cli_CallsRun(t *testing.T) {
@@ -20,7 +27,7 @@ func TestEnvironment_Cli_CallsRun(t *testing.T) {
 func TestEnvironment_DefaultCli_Empty(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
-	defaultEnv := NewEnvironment()
+	defaultEnv := testEnvironment()
 
 	assert.Equal(defaultEnv.Cli([]string{}), 1, "CLI with no args")
 }
@@ -28,7 +35,7 @@ func TestEnvironment_DefaultCli_Empty(t *testing.T) {
 func TestEnvironment_DefaultCli_Help(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
-	defaultEnv := NewEnvironment()
+	defaultEnv := testEnvironment()
 
 	// Test the basic version options
 	assert.Equal(defaultEnv.Cli([]string{"--help"}), 1, "--help should print")
@@ -38,7 +45,7 @@ func TestEnvironment_DefaultCli_Help(t *testing.T) {
 func TestEnvironment_DefaultCli_Version(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
-	defaultEnv := NewEnvironment()
+	defaultEnv := testEnvironment()
 
 	// Test the basic version options
 	assert.Equal(defaultEnv.Cli([]string{"version"}), 0, "version should work")
@@ -56,7 +63,7 @@ func TestEnvironment_DefaultCli_Version(t *testing.T) {
 func TestEnvironment_DefaultUi(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
-	defaultEnv := NewEnvironment()
+	defaultEnv := NewEnvironment(nil)
 	assert.NotNil(defaultEnv.Ui(), "default UI should not be nil")
 
 	rwUi, ok := defaultEnv.Ui().(*ReaderWriterUi)
@@ -67,5 +74,18 @@ func TestEnvironment_DefaultUi(t *testing.T) {
 
 func TestEnvironment_PrintHelp(t *testing.T) {
 	// Just call the function and verify that no panics occur
-	NewEnvironment().PrintHelp()
+	testEnvironment().PrintHelp()
+}
+
+func TestEnvironment_SettingUi(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	ui := &ReaderWriterUi{new(bytes.Buffer), new(bytes.Buffer)}
+
+	config := &EnvironmentConfig{}
+	config.ui = ui
+
+	env := NewEnvironment(config)
+
+	assert.Equal(env.Ui(), ui, "UIs should be equal")
 }