From 4832244709d8a09c5199121d530f331d8e5bfa6d Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Sat, 29 Aug 2015 14:22:24 -0700
Subject: [PATCH] client: test updating alloc status

---
 client/client.go      | 15 +++++++++++++++
 client/client_test.go | 36 ++++++++++++++++++++++++++++++++++++
 nomad/server.go       |  7 +++++++
 3 files changed, 58 insertions(+)

diff --git a/client/client.go b/client/client.go
index 64cb706a5..cbb242335 100644
--- a/client/client.go
+++ b/client/client.go
@@ -406,6 +406,21 @@ func (c *Client) updateNodeStatus() error {
 	return nil
 }
 
+// updateAllocStatus is used to update the status of an allocation
+func (c *Client) updateAllocStatus(alloc *structs.Allocation) error {
+	args := structs.AllocUpdateRequest{
+		Alloc:        []*structs.Allocation{alloc},
+		WriteRequest: structs.WriteRequest{Region: c.config.Region},
+	}
+	var resp structs.GenericResponse
+	err := c.RPC("Client.UpdateAlloc", &args, &resp)
+	if err != nil {
+		c.logger.Printf("[ERR] client: failed to update allocation: %v", err)
+		return err
+	}
+	return nil
+}
+
 // watchAllocations is used to scan for updates to allocations
 func (c *Client) watchAllocations(allocUpdates chan []*structs.Allocation) {
 	req := structs.NodeSpecificRequest{
diff --git a/client/client_test.go b/client/client_test.go
index 4a7789d4b..e05184106 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/nomad"
+	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/hashicorp/nomad/testutil"
 )
@@ -202,3 +203,38 @@ func TestClient_Heartbeat(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	})
 }
+
+func TestClient_UpdateAllocStatus(t *testing.T) {
+	s1, _ := testServer(t, nil)
+	defer s1.Shutdown()
+	testutil.WaitForLeader(t, s1.RPC)
+
+	c1 := testClient(t, func(c *config.Config) {
+		c.RPCHandler = s1
+	})
+	defer c1.Shutdown()
+
+	alloc := mock.Alloc()
+	alloc.NodeID = c1.Node().ID
+
+	state := s1.State()
+	state.UpdateAllocations(100, []*structs.Allocation{alloc})
+
+	newAlloc := new(structs.Allocation)
+	*newAlloc = *alloc
+	newAlloc.ClientStatus = structs.AllocClientStatusRunning
+
+	err := c1.updateAllocStatus(newAlloc)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	out, err := state.GetAllocByID(alloc.ID)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if out == nil || out.ClientStatus != structs.AllocClientStatusRunning {
+		t.Fatalf("bad: %#v", out)
+	}
+}
diff --git a/nomad/server.go b/nomad/server.go
index 1aaa58844..a70628ec1 100644
--- a/nomad/server.go
+++ b/nomad/server.go
@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/hashicorp/consul/tlsutil"
+	"github.com/hashicorp/nomad/nomad/state"
 	"github.com/hashicorp/raft"
 	"github.com/hashicorp/raft-boltdb"
 	"github.com/hashicorp/serf/serf"
@@ -590,6 +591,12 @@ func (s *Server) Encrypted() bool {
 	return s.serf.EncryptionEnabled()
 }
 
+// State returns the underlying state store. This should *not*
+// be used to modify state directly.
+func (s *Server) State() *state.StateStore {
+	return s.fsm.State()
+}
+
 // inmemCodec is used to do an RPC call without going over a network
 type inmemCodec struct {
 	method string