client: test updating alloc status

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{

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)
+	}
+}

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