fixup: Validate restart count

nomad/client_alloc_endpoint_test.go

@@ -861,11 +861,35 @@ func TestClientAllocations_Restart_Local(t *testing.T) {
 	require.NotNil(err)
 	require.Contains(err.Error(), "missing")
 
-	// Fetch the response setting the node id
+	// Fetch the response setting the alloc id - This should not error becuase the
+	// alloc is running.
 	req.AllocID = a.ID
 	var resp2 structs.GenericResponse
 	err = msgpackrpc.CallWithCodec(codec, "ClientAllocations.Restart", req, &resp2)
 	require.Nil(err)
+
+	testutil.WaitForResult(func() (bool, error) {
+		alloc, err := state.AllocByID(nil, a.ID)
+		if err != nil {
+			return false, err
+		}
+		if alloc == nil {
+			return false, fmt.Errorf("unknown alloc")
+		}
+
+		taskState := alloc.TaskStates["web"]
+		if taskState == nil {
+			return false, fmt.Errorf("could not find task state")
+		}
+
+		if taskState.Restarts != 1 {
+			return false, fmt.Errorf("expected task 'web' to have 1 restart, got: %d", taskState.Restarts)
+		}
+
+		return true, nil
+	}, func(err error) {
+		t.Fatalf("Alloc on node %q not running: %v", c.NodeID(), err)
+	})
 }
 
 func TestClientAllocations_Restart_Remote(t *testing.T) {
@@ -951,7 +975,8 @@ func TestClientAllocations_Restart_Remote(t *testing.T) {
 	require.NotNil(err)
 	require.Contains(err.Error(), "missing")
 
-	// Fetch the response setting the node id
+	// Fetch the response setting the alloc id - This should succeed because the
+	// alloc is running
 	req.AllocID = a.ID
 	var resp2 structs.GenericResponse
 	err = msgpackrpc.CallWithCodec(codec, "ClientAllocations.Restart", req, &resp2)