From 65e79fb7333431ed65ab594f19dbf9755bdeb956 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sun, 30 Aug 2015 16:35:04 -0700
Subject: [PATCH] client: refactor AllocRunner for testing

---
 client/alloc_runner.go      | 17 ++++++++-------
 client/alloc_runner_test.go | 42 +++++++++++++++++++++++++++++++++++++
 client/client.go            |  4 ++--
 3 files changed, 54 insertions(+), 9 deletions(-)
 create mode 100644 client/alloc_runner_test.go

diff --git a/client/alloc_runner.go b/client/alloc_runner.go
index 9575b31e6..69734df50 100644
--- a/client/alloc_runner.go
+++ b/client/alloc_runner.go
@@ -27,11 +27,14 @@ type taskStatus struct {
 	Description string
 }
 
+// AllocStateUpdater is used to update the status of an allocation
+type AllocStateUpdater func(alloc *structs.Allocation) error
+
 // AllocRunner is used to wrap an allocation and provide the execution context.
 type AllocRunner struct {
-	config *config.Config
-	client *Client
-	logger *log.Logger
+	config  *config.Config
+	updater AllocStateUpdater
+	logger  *log.Logger
 
 	alloc *structs.Allocation
 
@@ -59,11 +62,11 @@ type allocRunnerState struct {
 }
 
 // NewAllocRunner is used to create a new allocation context
-func NewAllocRunner(config *config.Config, client *Client, alloc *structs.Allocation) *AllocRunner {
+func NewAllocRunner(logger *log.Logger, config *config.Config, updater AllocStateUpdater, alloc *structs.Allocation) *AllocRunner {
 	ar := &AllocRunner{
 		config:     config,
-		client:     client,
-		logger:     client.logger,
+		updater:    updater,
+		logger:     logger,
 		alloc:      alloc,
 		dirtyCh:    make(chan struct{}, 1),
 		tasks:      make(map[string]*TaskRunner),
@@ -197,7 +200,7 @@ func (r *AllocRunner) syncStatus() {
 		}
 
 		// Attempt to update the status
-		if err := r.client.updateAllocStatus(r.alloc); err != nil {
+		if err := r.updater(r.alloc); err != nil {
 			r.logger.Printf("[ERR] client: failed to update alloc '%s' status to %s: %s",
 				r.alloc.ID, r.alloc.ClientStatus, err)
 			retryCh = time.After(allocSyncRetryIntv + randomStagger(allocSyncRetryIntv))
diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
new file mode 100644
index 000000000..bd605c3b1
--- /dev/null
+++ b/client/alloc_runner_test.go
@@ -0,0 +1,42 @@
+package client
+
+import (
+	"testing"
+
+	"github.com/hashicorp/nomad/nomad/mock"
+	"github.com/hashicorp/nomad/nomad/structs"
+)
+
+type MockAllocStateUpdater struct {
+	Count  int
+	Allocs []*structs.Allocation
+	Err    error
+}
+
+func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) error {
+	m.Count += 1
+	m.Allocs = append(m.Allocs, alloc)
+	return m.Err
+}
+
+func testAllocRunner() (*MockAllocStateUpdater, *AllocRunner) {
+	logger := testLogger()
+	conf := DefaultConfig()
+	conf.StateDir = "/tmp"
+	upd := &MockAllocStateUpdater{}
+	alloc := mock.Alloc()
+	ar := NewAllocRunner(logger, conf, upd.Update, alloc)
+	return upd, ar
+}
+
+func TestAllocRunner_SimpleRun(t *testing.T) {
+}
+
+func TestAllocRunner_Update(t *testing.T) {
+}
+
+func TestAllocRunner_Destroy(t *testing.T) {
+}
+
+func TestAllocRunner_SaveRestoreState(t *testing.T) {
+}
diff --git a/client/client.go b/client/client.go
index ce9f93fb5..8bde19685 100644
--- a/client/client.go
+++ b/client/client.go
@@ -236,7 +236,7 @@ func (c *Client) restoreState() error {
 	for _, entry := range list {
 		id := entry.Name()
 		alloc := &structs.Allocation{ID: id}
-		ar := NewAllocRunner(c.config, c, alloc)
+		ar := NewAllocRunner(c.logger, c.config, c.updateAllocStatus, alloc)
 		c.allocs[id] = ar
 		if err := ar.RestoreState(); err != nil {
 			c.logger.Printf("[ERR] client: failed to restore state for alloc %s: %v",
@@ -576,7 +576,7 @@ func (c *Client) updateAlloc(exist, update *structs.Allocation) error {
 func (c *Client) addAlloc(alloc *structs.Allocation) error {
 	c.allocLock.Lock()
 	defer c.allocLock.Unlock()
-	ar := NewAllocRunner(c.config, c, alloc)
+	ar := NewAllocRunner(c.logger, c.config, c.updateAllocStatus, alloc)
 	c.allocs[alloc.ID] = ar
 	go ar.Run()
 	return nil