packer/rpc: Test BuilderFactory result works. Close back RPC after one

conn

packer/rpc/builder_factory.go

@@ -3,7 +3,6 @@ package rpc
 import (
 	"github.com/mitchellh/packer/packer"
 	"net/rpc"
-	"time"
 )
 
 // An implementation of packer.BuilderFactory where the factory is actually
@@ -39,13 +38,7 @@ func (b *BuilderFactoryServer) CreateBuilder(args *BuilderFactoryCreateArgs, rep
 	// Now we wrap that back up into a server, and send it on backwards.
 	server := NewServer()
 	server.RegisterBuilder(builder)
-	server.Start()
-
-	// Start a timer where we'll quit the server after some number of seconds.
-	go func() {
-		time.Sleep(30 * time.Second)
-		server.Stop()
-	}()
+	server.StartSingle()
 
 	// Set the reply to the address of the sever
 	*reply = server.Address()

packer/rpc/builder_factory_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 )
 
+var createResult = &testBuilder{}
+
 type testBuilderFactory struct {
 	createCalled bool
 	createName string
@@ -15,7 +17,7 @@ type testBuilderFactory struct {
 func (b *testBuilderFactory) CreateBuilder(name string) packer.Builder {
 	b.createCalled = true
 	b.createName = name
-	return &testBuilder{}
+	return createResult
 }
 
 func TestBuilderFactoryRPC(t *testing.T) {
@@ -37,9 +39,13 @@ func TestBuilderFactoryRPC(t *testing.T) {
 	// Test Create
 	name := "foo"
 	bClient := &BuilderFactory{client}
-	_ = bClient.CreateBuilder(name)
+	builder := bClient.CreateBuilder(name)
 	assert.True(b.createCalled, "create should be called")
 	assert.Equal(b.createName, "foo", "name should be foo")
+
+	builder.Prepare(42)
+	assert.True(createResult.prepareCalled, "prepare should be called")
+	assert.Equal(createResult.prepareConfig, 42, "42 should be config")
 }
 
 func TestBuilderFactory_ImplementsBuilderFactory(t *testing.T) {

packer/rpc/server.go

@@ -46,6 +46,21 @@ func (s *Server) RegisterUi(ui packer.Ui) {
 }
 
 func (s *Server) Start() error {
+	return s.start(false)
+}
+
+func (s *Server) StartSingle() error {
+	return s.start(true)
+}
+
+func (s *Server) Stop() {
+	if s.listener != nil {
+		s.listener.Close()
+		s.listener = nil
+	}
+}
+
+func (s *Server) start(singleConn bool) error {
 	if s.listener != nil {
 		return errors.New("Server already started.")
 	}
@@ -66,15 +81,15 @@ func (s *Server) Start() error {
 			}
 
 			go s.server.ServeConn(conn)
+
+			// If we're only accepting a single connection then
+			// stop.
+			if singleConn {
+				s.Stop()
+				break
+			}
 		}
 	}(s.listener)
 
 	return nil
 }
-
-func (s *Server) Stop() {
-	if s.listener != nil {
-		s.listener.Close()
-		s.listener = nil
-	}
-}