From 86f1fbe9256158b116afea07096792921ce5b63e Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 5 May 2013 15:12:55 -0700
Subject: [PATCH] packer/rpc: Support the new Builder func on Environment

---
 packer/environment.go          |  7 +++++++
 packer/environment_test.go     | 14 ++++++++++++++
 packer/rpc/environment.go      | 21 +++++++++++++++++++++
 packer/rpc/environment_test.go | 18 +++++++++++++++++-
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/packer/environment.go b/packer/environment.go
index 5f6e7681d..b3e6a6a6a 100644
--- a/packer/environment.go
+++ b/packer/environment.go
@@ -19,6 +19,7 @@ type CommandFunc func(name string) Command
 // It allows for things such as executing CLI commands, getting the
 // list of available builders, and more.
 type Environment interface {
+	Builder(name string) Builder
 	Cli(args []string) int
 	Ui() Ui
 }
@@ -68,6 +69,12 @@ func NewEnvironment(config *EnvironmentConfig) (resultEnv Environment, err error
 	return
 }
 
+// Returns a builder of the given name that is registered with this
+// environment.
+func (e *coreEnvironment) Builder(name string) Builder {
+	return e.builderFunc(name)
+}
+
 // Executes a command as if it was typed on the command-line interface.
 // The return value is the exit code of the command.
 func (e *coreEnvironment) Cli(args []string) int {
diff --git a/packer/environment_test.go b/packer/environment_test.go
index 0a28cc988..3d570d820 100644
--- a/packer/environment_test.go
+++ b/packer/environment_test.go
@@ -50,6 +50,20 @@ func TestNewEnvironment_NoConfig(t *testing.T) {
 	assert.NotNil(err, "should be an error")
 }
 
+func TestEnvironment_Builder(t *testing.T) {
+	assert := asserts.NewTestingAsserts(t, true)
+
+	builder := &TestBuilder{}
+	builders := make(map[string]Builder)
+	builders["foo"] = builder
+
+	config := DefaultEnvironmentConfig()
+	config.BuilderFunc = func(n string) Builder { return builders[n] }
+
+	env, _ := NewEnvironment(config)
+	assert.Equal(env.Builder("foo"), builder, "should return correct builder")
+}
+
 func TestEnvironment_Cli_CallsRun(t *testing.T) {
 	assert := asserts.NewTestingAsserts(t, true)
 
diff --git a/packer/rpc/environment.go b/packer/rpc/environment.go
index c05f07832..236bb1c95 100644
--- a/packer/rpc/environment.go
+++ b/packer/rpc/environment.go
@@ -21,6 +21,15 @@ type EnvironmentCliArgs struct {
 	Args []string
 }
 
+func (e *Environment) Builder(name string) packer.Builder {
+	var reply string
+	e.client.Call("Environment.Builder", name, &reply)
+
+	// TODO: error handling
+	client, _ := rpc.Dial("tcp", reply)
+	return &Builder{client}
+}
+
 func (e *Environment) Cli(args []string) (result int) {
 	rpcArgs := &EnvironmentCliArgs{args}
 	e.client.Call("Environment.Cli", rpcArgs, &result)
@@ -36,6 +45,18 @@ func (e *Environment) Ui() packer.Ui {
 	return &Ui{client}
 }
 
+func (e *EnvironmentServer) Builder(name *string, reply *string) error {
+	builder := e.env.Builder(*name)
+
+	// Wrap it
+	server := NewServer()
+	server.RegisterBuilder(builder)
+	server.StartSingle()
+
+	*reply = server.Address()
+	return nil
+}
+
 func (e *EnvironmentServer) Cli(args *EnvironmentCliArgs, reply *int) error {
 	*reply = e.env.Cli(args.Args)
 	return nil
diff --git a/packer/rpc/environment_test.go b/packer/rpc/environment_test.go
index 6e0541220..6699bfc26 100644
--- a/packer/rpc/environment_test.go
+++ b/packer/rpc/environment_test.go
@@ -7,15 +7,23 @@ import (
 	"testing"
 )
 
+var testEnvBuilder = &testBuilder{}
 var testEnvUi = &testUi{}
 
 type testEnvironment struct {
-	bfCalled bool
+	builderCalled bool
+	builderName string
 	cliCalled bool
 	cliArgs []string
 	uiCalled bool
 }
 
+func (e *testEnvironment) Builder(name string) packer.Builder {
+	e.builderCalled = true
+	e.builderName = name
+	return testEnvBuilder
+}
+
 func (e *testEnvironment) Cli(args []string) int {
 	e.cliCalled = true
 	e.cliArgs = args
@@ -44,6 +52,14 @@ func TestEnvironmentRPC(t *testing.T) {
 	assert.Nil(err, "should be able to connect")
 	eClient := &Environment{client}
 
+	// Test Builder
+	builder := eClient.Builder("foo")
+	assert.True(e.builderCalled, "Builder should be called")
+	assert.Equal(e.builderName, "foo", "Correct name for Builder")
+
+	builder.Prepare(nil)
+	assert.True(testEnvBuilder.prepareCalled, "Prepare should be called")
+
 	// Test Cli
 	cliArgs := []string{"foo", "bar"}
 	result := eClient.Cli(cliArgs)