From ac83cf652a7bcce96be82c6981bdb9db93a926a3 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Wed, 8 May 2013 13:11:47 -0700
Subject: [PATCH] packer/rpc: Setup the Environment properly for Command

---
 packer/rpc/command.go      | 18 ++++++++++++++----
 packer/rpc/command_test.go | 13 +++++++++++--
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/packer/rpc/command.go b/packer/rpc/command.go
index 0e2c89e39..b1662a597 100644
--- a/packer/rpc/command.go
+++ b/packer/rpc/command.go
@@ -18,7 +18,7 @@ type ServerCommand struct {
 }
 
 type CommandRunArgs struct {
-	Env packer.Environment
+	RPCAddress string
 	Args []string
 }
 
@@ -29,8 +29,11 @@ func Command(client *rpc.Client) *ClientCommand {
 }
 
 func (c *ClientCommand) Run(env packer.Environment, args []string) (result int) {
-	// TODO: Environment
-	rpcArgs := &CommandRunArgs{nil, args}
+	// Create and start the server for the Environment
+	server := rpc.NewServer()
+	RegisterEnvironment(server, env)
+
+	rpcArgs := &CommandRunArgs{serveSingleConn(server), args}
 	err := c.client.Call("Command.Run", rpcArgs, &result)
 	if err != nil {
 		panic(err)
@@ -49,7 +52,14 @@ func (c *ClientCommand) Synopsis() (result string) {
 }
 
 func (c *ServerCommand) Run(args *CommandRunArgs, reply *int) error {
-	*reply = c.command.Run(args.Env, args.Args)
+	client, err := rpc.Dial("tcp", args.RPCAddress)
+	if err != nil {
+		return err
+	}
+
+	env := &Environment{client}
+
+	*reply = c.command.Run(env, args.Args)
 	return nil
 }
 
diff --git a/packer/rpc/command_test.go b/packer/rpc/command_test.go
index 238e163d6..d3a4f3904 100644
--- a/packer/rpc/command_test.go
+++ b/packer/rpc/command_test.go
@@ -85,12 +85,21 @@ func TestRPCCommand(t *testing.T) {
 	}
 
 	clientComm := &ClientCommand{client}
+
+	// Test run
 	runArgs := []string{"foo", "bar"}
 	testEnv := &testEnvironment{}
 	exitCode := clientComm.Run(testEnv, runArgs)
-	synopsis := clientComm.Synopsis()
-
 	assert.Equal(command.runArgs, runArgs, "Correct args should be sent")
 	assert.Equal(exitCode, 0, "Exit code should be correct")
+
+	assert.NotNil(command.runEnv, "should have an env")
+	if command.runEnv != nil {
+		command.runEnv.Ui()
+		assert.True(testEnv.uiCalled, "UI should be called on env")
+	}
+
+	// Test Synopsis
+	synopsis := clientComm.Synopsis()
 	assert.Equal(synopsis, "foo", "Synopsis should be correct")
 }