open-nomad/nomad/state/state_store_test.go

package state

import (
	"context"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/go-memdb"
	"github.com/kr/pretty"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/hashicorp/nomad/helper"
	"github.com/hashicorp/nomad/helper/uuid"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/structs"
)

func testStateStore(t *testing.T) *StateStore {
	return TestStateStore(t)
}

func TestStateStore_Blocking_Error(t *testing.T) {
	t.Parallel()

	expected := fmt.Errorf("test error")
	errFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
		return nil, 0, expected
	}

	state := testStateStore(t)
	_, idx, err := state.BlockingQuery(errFn, 10, context.Background())
	assert.EqualError(t, err, expected.Error())
	assert.Zero(t, idx)
}

func TestStateStore_Blocking_Timeout(t *testing.T) {
	t.Parallel()

	noopFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
		return nil, 5, nil
	}

	state := testStateStore(t)
	timeout := time.Now().Add(250 * time.Millisecond)
	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
	defer cancel()

	_, idx, err := state.BlockingQuery(noopFn, 10, deadlineCtx)
	assert.EqualError(t, err, context.DeadlineExceeded.Error())
	assert.EqualValues(t, 5, idx)
	assert.WithinDuration(t, timeout, time.Now(), 100*time.Millisecond)
}

func TestStateStore_Blocking_MinQuery(t *testing.T) {
	t.Parallel()

	node := mock.Node()
	count := 0
	queryFn := func(ws memdb.WatchSet, s *StateStore) (interface{}, uint64, error) {
		_, err := s.NodeByID(ws, node.ID)
		if err != nil {
			return nil, 0, err
		}

		count++
		if count == 1 {
			return false, 5, nil
		} else if count > 2 {
			return false, 20, fmt.Errorf("called too many times")
		}

		return true, 11, nil
	}

	state := testStateStore(t)
	timeout := time.Now().Add(100 * time.Millisecond)
	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
	defer cancel()

	time.AfterFunc(5*time.Millisecond, func() {
		state.UpsertNode(structs.MsgTypeTestSetup, 11, node)
	})

	resp, idx, err := state.BlockingQuery(queryFn, 10, deadlineCtx)
	if assert.Nil(t, err) {
		assert.Equal(t, 2, count)
		assert.EqualValues(t, 11, idx)
		assert.True(t, resp.(bool))
	}
}

// COMPAT 0.11: Uses AllocUpdateRequest.Alloc
// This test checks that:
// 1) The job is denormalized
// 2) Allocations are created
func TestStateStore_UpsertPlanResults_AllocationsCreated_Denormalized(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	job := alloc.Job
	alloc.Job = nil

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	eval := mock.Eval()
	eval.JobID = job.ID

	// Create an eval
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a plan result
	res := structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			Alloc: []*structs.Allocation{alloc},
			Job:   job,
		},
		EvalID: eval.ID,
	}
	assert := assert.New(t)
	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
	assert.Nil(err)

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	assert.Nil(err)
	assert.Equal(alloc, out)

	index, err := state.Index("allocs")
	assert.Nil(err)
	assert.EqualValues(1000, index)

	if watchFired(ws) {
		t.Fatalf("bad")
	}

	evalOut, err := state.EvalByID(ws, eval.ID)
	assert.Nil(err)
	assert.NotNil(evalOut)
	assert.EqualValues(1000, evalOut.ModifyIndex)
}

// This test checks that:
// 1) The job is denormalized
// 2) Allocations are denormalized and updated with the diff
// That stopped allocs Job is unmodified
func TestStateStore_UpsertPlanResults_AllocationsDenormalized(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	job := alloc.Job
	alloc.Job = nil

	stoppedAlloc := mock.Alloc()
	stoppedAlloc.Job = job
	stoppedAllocDiff := &structs.AllocationDiff{
		ID:                 stoppedAlloc.ID,
		DesiredDescription: "desired desc",
		ClientStatus:       structs.AllocClientStatusLost,
	}
	preemptedAlloc := mock.Alloc()
	preemptedAlloc.Job = job
	preemptedAllocDiff := &structs.AllocationDiff{
		ID:                    preemptedAlloc.ID,
		PreemptedByAllocation: alloc.ID,
	}

	require := require.New(t)
	require.NoError(state.UpsertAllocs(structs.MsgTypeTestSetup, 900, []*structs.Allocation{stoppedAlloc, preemptedAlloc}))
	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 999, job))

	// modify job and ensure that stopped and preempted alloc point to original Job
	mJob := job.Copy()
	mJob.TaskGroups[0].Name = "other"

	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1001, mJob))

	eval := mock.Eval()
	eval.JobID = job.ID

	// Create an eval
	require.NoError(state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}))

	// Create a plan result
	res := structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			AllocsUpdated: []*structs.Allocation{alloc},
			AllocsStopped: []*structs.AllocationDiff{stoppedAllocDiff},
			Job:           mJob,
		},
		EvalID:          eval.ID,
		AllocsPreempted: []*structs.AllocationDiff{preemptedAllocDiff},
	}
	assert := assert.New(t)
	planModifyIndex := uint64(1000)
	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, planModifyIndex, &res)
	require.NoError(err)

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	require.NoError(err)
	assert.Equal(alloc, out)

	outJob, err := state.JobByID(ws, job.Namespace, job.ID)
	require.NoError(err)
	require.Equal(mJob.TaskGroups, outJob.TaskGroups)
	require.NotEmpty(job.TaskGroups, outJob.TaskGroups)

	updatedStoppedAlloc, err := state.AllocByID(ws, stoppedAlloc.ID)
	require.NoError(err)
	assert.Equal(stoppedAllocDiff.DesiredDescription, updatedStoppedAlloc.DesiredDescription)
	assert.Equal(structs.AllocDesiredStatusStop, updatedStoppedAlloc.DesiredStatus)
	assert.Equal(stoppedAllocDiff.ClientStatus, updatedStoppedAlloc.ClientStatus)
	assert.Equal(planModifyIndex, updatedStoppedAlloc.AllocModifyIndex)
	assert.Equal(planModifyIndex, updatedStoppedAlloc.AllocModifyIndex)
	assert.Equal(job.TaskGroups, updatedStoppedAlloc.Job.TaskGroups)

	updatedPreemptedAlloc, err := state.AllocByID(ws, preemptedAlloc.ID)
	require.NoError(err)
	assert.Equal(structs.AllocDesiredStatusEvict, updatedPreemptedAlloc.DesiredStatus)
	assert.Equal(preemptedAllocDiff.PreemptedByAllocation, updatedPreemptedAlloc.PreemptedByAllocation)
	assert.Equal(planModifyIndex, updatedPreemptedAlloc.AllocModifyIndex)
	assert.Equal(planModifyIndex, updatedPreemptedAlloc.AllocModifyIndex)
	assert.Equal(job.TaskGroups, updatedPreemptedAlloc.Job.TaskGroups)

	index, err := state.Index("allocs")
	require.NoError(err)
	assert.EqualValues(planModifyIndex, index)

	require.False(watchFired(ws))

	evalOut, err := state.EvalByID(ws, eval.ID)
	require.NoError(err)
	require.NotNil(evalOut)
	assert.EqualValues(planModifyIndex, evalOut.ModifyIndex)

}

// This test checks that the deployment is created and allocations count towards
// the deployment
func TestStateStore_UpsertPlanResults_Deployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	alloc2 := mock.Alloc()
	job := alloc.Job
	alloc.Job = nil
	alloc2.Job = nil

	d := mock.Deployment()
	alloc.DeploymentID = d.ID
	alloc2.DeploymentID = d.ID

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	eval := mock.Eval()
	eval.JobID = job.ID

	// Create an eval
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a plan result
	res := structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			Alloc: []*structs.Allocation{alloc, alloc2},
			Job:   job,
		},
		Deployment: d,
		EvalID:     eval.ID,
	}

	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	assert := assert.New(t)
	out, err := state.AllocByID(ws, alloc.ID)
	assert.Nil(err)
	assert.Equal(alloc, out)

	dout, err := state.DeploymentByID(ws, d.ID)
	assert.Nil(err)
	assert.NotNil(dout)

	tg, ok := dout.TaskGroups[alloc.TaskGroup]
	assert.True(ok)
	assert.NotNil(tg)
	assert.Equal(2, tg.PlacedAllocs)

	evalOut, err := state.EvalByID(ws, eval.ID)
	assert.Nil(err)
	assert.NotNil(evalOut)
	assert.EqualValues(1000, evalOut.ModifyIndex)

	if watchFired(ws) {
		t.Fatalf("bad")
	}

	// Update the allocs to be part of a new deployment
	d2 := d.Copy()
	d2.ID = uuid.Generate()

	allocNew := alloc.Copy()
	allocNew.DeploymentID = d2.ID
	allocNew2 := alloc2.Copy()
	allocNew2.DeploymentID = d2.ID

	// Create another plan
	res = structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			Alloc: []*structs.Allocation{allocNew, allocNew2},
			Job:   job,
		},
		Deployment: d2,
		EvalID:     eval.ID,
	}

	err = state.UpsertPlanResults(structs.MsgTypeTestSetup, 1001, &res)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	dout, err = state.DeploymentByID(ws, d2.ID)
	assert.Nil(err)
	assert.NotNil(dout)

	tg, ok = dout.TaskGroups[alloc.TaskGroup]
	assert.True(ok)
	assert.NotNil(tg)
	assert.Equal(2, tg.PlacedAllocs)

	evalOut, err = state.EvalByID(ws, eval.ID)
	assert.Nil(err)
	assert.NotNil(evalOut)
	assert.EqualValues(1001, evalOut.ModifyIndex)
}

// This test checks that:
// 1) Preempted allocations in plan results are updated
// 2) Evals are inserted for preempted jobs
func TestStateStore_UpsertPlanResults_PreemptedAllocs(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	alloc := mock.Alloc()
	job := alloc.Job
	alloc.Job = nil

	// Insert job
	err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job)
	require.NoError(err)

	// Create an eval
	eval := mock.Eval()
	eval.JobID = job.ID
	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval})
	require.NoError(err)

	// Insert alloc that will be preempted in the plan
	preemptedAlloc := mock.Alloc()
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 2, []*structs.Allocation{preemptedAlloc})
	require.NoError(err)

	minimalPreemptedAlloc := &structs.Allocation{
		ID:                    preemptedAlloc.ID,
		PreemptedByAllocation: alloc.ID,
		ModifyTime:            time.Now().Unix(),
	}

	// Create eval for preempted job
	eval2 := mock.Eval()
	eval2.JobID = preemptedAlloc.JobID

	// Create a plan result
	res := structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			Alloc: []*structs.Allocation{alloc},
			Job:   job,
		},
		EvalID:          eval.ID,
		NodePreemptions: []*structs.Allocation{minimalPreemptedAlloc},
		PreemptionEvals: []*structs.Evaluation{eval2},
	}

	err = state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
	require.NoError(err)

	ws := memdb.NewWatchSet()

	// Verify alloc and eval created by plan
	out, err := state.AllocByID(ws, alloc.ID)
	require.NoError(err)
	require.Equal(alloc, out)

	index, err := state.Index("allocs")
	require.NoError(err)
	require.EqualValues(1000, index)

	evalOut, err := state.EvalByID(ws, eval.ID)
	require.NoError(err)
	require.NotNil(evalOut)
	require.EqualValues(1000, evalOut.ModifyIndex)

	// Verify preempted alloc
	preempted, err := state.AllocByID(ws, preemptedAlloc.ID)
	require.NoError(err)
	require.Equal(preempted.DesiredStatus, structs.AllocDesiredStatusEvict)
	require.Equal(preempted.DesiredDescription, fmt.Sprintf("Preempted by alloc ID %v", alloc.ID))
	require.Equal(preempted.Job.ID, preemptedAlloc.Job.ID)
	require.Equal(preempted.Job, preemptedAlloc.Job)

	// Verify eval for preempted job
	preemptedJobEval, err := state.EvalByID(ws, eval2.ID)
	require.NoError(err)
	require.NotNil(preemptedJobEval)
	require.EqualValues(1000, preemptedJobEval.ModifyIndex)

}

// This test checks that deployment updates are applied correctly
func TestStateStore_UpsertPlanResults_DeploymentUpdates(t *testing.T) {
	t.Parallel()
	state := testStateStore(t)

	// Create a job that applies to all
	job := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a deployment that we will update its status
	doutstanding := mock.Deployment()
	doutstanding.JobID = job.ID

	if err := state.UpsertDeployment(1000, doutstanding); err != nil {
		t.Fatalf("err: %v", err)
	}

	eval := mock.Eval()
	eval.JobID = job.ID

	// Create an eval
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}
	alloc := mock.Alloc()
	alloc.Job = nil

	dnew := mock.Deployment()
	dnew.JobID = job.ID
	alloc.DeploymentID = dnew.ID

	// Update the old deployment
	update := &structs.DeploymentStatusUpdate{
		DeploymentID:      doutstanding.ID,
		Status:            "foo",
		StatusDescription: "bar",
	}

	// Create a plan result
	res := structs.ApplyPlanResultsRequest{
		AllocUpdateRequest: structs.AllocUpdateRequest{
			Alloc: []*structs.Allocation{alloc},
			Job:   job,
		},
		Deployment:        dnew,
		DeploymentUpdates: []*structs.DeploymentStatusUpdate{update},
		EvalID:            eval.ID,
	}

	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	assert := assert.New(t)
	ws := memdb.NewWatchSet()

	// Check the deployments are correctly updated.
	dout, err := state.DeploymentByID(ws, dnew.ID)
	assert.Nil(err)
	assert.NotNil(dout)

	tg, ok := dout.TaskGroups[alloc.TaskGroup]
	assert.True(ok)
	assert.NotNil(tg)
	assert.Equal(1, tg.PlacedAllocs)

	doutstandingout, err := state.DeploymentByID(ws, doutstanding.ID)
	assert.Nil(err)
	assert.NotNil(doutstandingout)
	assert.Equal(update.Status, doutstandingout.Status)
	assert.Equal(update.StatusDescription, doutstandingout.StatusDescription)
	assert.EqualValues(1000, doutstandingout.ModifyIndex)

	evalOut, err := state.EvalByID(ws, eval.ID)
	assert.Nil(err)
	assert.NotNil(evalOut)
	assert.EqualValues(1000, evalOut.ModifyIndex)
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertDeployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	deployment := mock.Deployment()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.DeploymentsByJobID(ws, deployment.Namespace, deployment.ID, true)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.UpsertDeployment(1000, deployment)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.DeploymentByID(ws, deployment.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(deployment, out) {
		t.Fatalf("bad: %#v %#v", deployment, out)
	}

	index, err := state.Index("deployment")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

// Tests that deployments of older create index and same job id are not returned
func TestStateStore_OldDeployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	job.ID = "job1"
	state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)

	deploy1 := mock.Deployment()
	deploy1.JobID = job.ID
	deploy1.JobCreateIndex = job.CreateIndex

	deploy2 := mock.Deployment()
	deploy2.JobID = job.ID
	deploy2.JobCreateIndex = 11

	require := require.New(t)

	// Insert both deployments
	err := state.UpsertDeployment(1001, deploy1)
	require.Nil(err)

	err = state.UpsertDeployment(1002, deploy2)
	require.Nil(err)

	ws := memdb.NewWatchSet()
	// Should return both deployments
	deploys, err := state.DeploymentsByJobID(ws, deploy1.Namespace, job.ID, true)
	require.Nil(err)
	require.Len(deploys, 2)

	// Should only return deploy1
	deploys, err = state.DeploymentsByJobID(ws, deploy1.Namespace, job.ID, false)
	require.Nil(err)
	require.Len(deploys, 1)
	require.Equal(deploy1.ID, deploys[0].ID)
}

func TestStateStore_DeleteDeployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	d1 := mock.Deployment()
	d2 := mock.Deployment()

	err := state.UpsertDeployment(1000, d1)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if err := state.UpsertDeployment(1001, d2); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.DeploymentByID(ws, d1.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.DeleteDeployment(1002, []string{d1.ID, d2.ID})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.DeploymentByID(ws, d1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", d1, out)
	}

	index, err := state.Index("deployment")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1002 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Deployments(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var deployments []*structs.Deployment

	for i := 0; i < 10; i++ {
		deployment := mock.Deployment()
		deployments = append(deployments, deployment)

		err := state.UpsertDeployment(1000+uint64(i), deployment)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.Deployments(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*structs.Deployment
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Deployment))
	}

	lessThan := func(i, j int) bool {
		return deployments[i].ID < deployments[j].ID
	}
	sort.Slice(deployments, lessThan)
	sort.Slice(out, lessThan)

	if !reflect.DeepEqual(deployments, out) {
		t.Fatalf("bad: %#v %#v", deployments, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeploymentsByIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	deploy := mock.Deployment()

	deploy.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
	err := state.UpsertDeployment(1000, deploy)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that getters don't cause it to fire
	ws := memdb.NewWatchSet()
	iter, err := state.DeploymentsByIDPrefix(ws, deploy.Namespace, deploy.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	gatherDeploys := func(iter memdb.ResultIterator) []*structs.Deployment {
		var deploys []*structs.Deployment
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			deploy := raw.(*structs.Deployment)
			deploys = append(deploys, deploy)
		}
		return deploys
	}

	deploys := gatherDeploys(iter)
	if len(deploys) != 1 {
		t.Fatalf("err: %v", err)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}

	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	deploys = gatherDeploys(iter)
	if len(deploys) != 1 {
		t.Fatalf("err: %v", err)
	}

	deploy = mock.Deployment()
	deploy.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
	err = state.UpsertDeployment(1001, deploy)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	deploys = gatherDeploys(iter)
	if len(deploys) != 2 {
		t.Fatalf("err: %v", err)
	}

	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "1111")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	deploys = gatherDeploys(iter)
	if len(deploys) != 1 {
		t.Fatalf("err: %v", err)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertNode_Node(t *testing.T) {
	t.Parallel()

	require := require.New(t)
	state := testStateStore(t)
	node := mock.Node()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.NodeByID(ws, node.ID)
	require.NoError(err)

	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))
	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.NoError(err)

	out2, err := state.NodeBySecretID(ws, node.SecretID)
	require.NoError(err)
	require.EqualValues(node, out)
	require.EqualValues(node, out2)
	require.Len(out.Events, 1)
	require.Equal(NodeRegisterEventRegistered, out.Events[0].Message)

	index, err := state.Index("nodes")
	require.NoError(err)
	require.EqualValues(1000, index)
	require.False(watchFired(ws))

	// Transition the node to down and then up and ensure we get a re-register
	// event
	down := out.Copy()
	down.Status = structs.NodeStatusDown
	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1001, down))
	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1002, out))

	out, err = state.NodeByID(ws, node.ID)
	require.NoError(err)
	require.Len(out.Events, 2)
	require.Equal(NodeRegisterEventReregistered, out.Events[1].Message)
}

func TestStateStore_DeleteNode_Node(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create and insert two nodes, which we'll delete
	node0 := mock.Node()
	node1 := mock.Node()
	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node0)
	require.NoError(t, err)
	err = state.UpsertNode(structs.MsgTypeTestSetup, 1001, node1)
	require.NoError(t, err)

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()

	// Check that both nodes are not nil
	out, err := state.NodeByID(ws, node0.ID)
	require.NoError(t, err)
	require.NotNil(t, out)
	out, err = state.NodeByID(ws, node1.ID)
	require.NoError(t, err)
	require.NotNil(t, out)

	// Delete both nodes in a batch, fires the watch
	err = state.DeleteNode(structs.MsgTypeTestSetup, 1002, []string{node0.ID, node1.ID})
	require.NoError(t, err)
	require.True(t, watchFired(ws))

	// Check that both nodes are nil
	ws = memdb.NewWatchSet()
	out, err = state.NodeByID(ws, node0.ID)
	require.NoError(t, err)
	require.Nil(t, out)
	out, err = state.NodeByID(ws, node1.ID)
	require.NoError(t, err)
	require.Nil(t, out)

	// Ensure that the index is still at 1002, from DeleteNode
	index, err := state.Index("nodes")
	require.NoError(t, err)
	require.Equal(t, uint64(1002), index)
	require.False(t, watchFired(ws))
}

func TestStateStore_UpdateNodeStatus_Node(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	node := mock.Node()

	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 800, node))

	// Create a watchset so we can test that update node status fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.NodeByID(ws, node.ID)
	require.NoError(err)

	event := &structs.NodeEvent{
		Message:   "Node ready foo",
		Subsystem: structs.NodeEventSubsystemCluster,
		Timestamp: time.Now(),
	}

	require.NoError(state.UpdateNodeStatus(structs.MsgTypeTestSetup, 801, node.ID, structs.NodeStatusReady, 70, event))
	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.NoError(err)
	require.Equal(structs.NodeStatusReady, out.Status)
	require.EqualValues(801, out.ModifyIndex)
	require.EqualValues(70, out.StatusUpdatedAt)
	require.Len(out.Events, 2)
	require.Equal(event.Message, out.Events[1].Message)

	index, err := state.Index("nodes")
	require.NoError(err)
	require.EqualValues(801, index)
	require.False(watchFired(ws))
}

func TestStateStore_BatchUpdateNodeDrain(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	n1, n2 := mock.Node(), mock.Node()
	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, n1))
	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1001, n2))

	// Create a watchset so we can test that update node drain fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.NodeByID(ws, n1.ID)
	require.Nil(err)

	expectedDrain := &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: -1 * time.Second,
		},
	}

	update := map[string]*structs.DrainUpdate{
		n1.ID: {
			DrainStrategy: expectedDrain,
		},
		n2.ID: {
			DrainStrategy: expectedDrain,
		},
	}

	event := &structs.NodeEvent{
		Message:   "Drain strategy enabled",
		Subsystem: structs.NodeEventSubsystemDrain,
		Timestamp: time.Now(),
	}
	events := map[string]*structs.NodeEvent{
		n1.ID: event,
		n2.ID: event,
	}

	require.Nil(state.BatchUpdateNodeDrain(structs.MsgTypeTestSetup, 1002, 7, update, events))
	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	for _, id := range []string{n1.ID, n2.ID} {
		out, err := state.NodeByID(ws, id)
		require.Nil(err)
		require.True(out.Drain)
		require.NotNil(out.DrainStrategy)
		require.Equal(out.DrainStrategy, expectedDrain)
		require.Len(out.Events, 2)
		require.EqualValues(1002, out.ModifyIndex)
		require.EqualValues(7, out.StatusUpdatedAt)
	}

	index, err := state.Index("nodes")
	require.Nil(err)
	require.EqualValues(1002, index)
	require.False(watchFired(ws))
}

func TestStateStore_UpdateNodeDrain_Node(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	node := mock.Node()

	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))

	// Create a watchset so we can test that update node drain fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.NodeByID(ws, node.ID)
	require.Nil(err)

	expectedDrain := &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: -1 * time.Second,
		},
	}

	event := &structs.NodeEvent{
		Message:   "Drain strategy enabled",
		Subsystem: structs.NodeEventSubsystemDrain,
		Timestamp: time.Now(),
	}
	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1001, node.ID, expectedDrain, false, 7, event))
	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.Nil(err)
	require.True(out.Drain)
	require.NotNil(out.DrainStrategy)
	require.Equal(out.DrainStrategy, expectedDrain)
	require.Len(out.Events, 2)
	require.EqualValues(1001, out.ModifyIndex)
	require.EqualValues(7, out.StatusUpdatedAt)

	index, err := state.Index("nodes")
	require.Nil(err)
	require.EqualValues(1001, index)
	require.False(watchFired(ws))
}

func TestStateStore_AddSingleNodeEvent(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	node := mock.Node()

	// We create a new node event every time we register a node
	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
	require.Nil(err)

	require.Equal(1, len(node.Events))
	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)

	// Create a watchset so we can test that AddNodeEvent fires the watch
	ws := memdb.NewWatchSet()
	_, err = state.NodeByID(ws, node.ID)
	require.Nil(err)

	nodeEvent := &structs.NodeEvent{
		Message:   "failed",
		Subsystem: "Driver",
		Timestamp: time.Now(),
	}
	nodeEvents := map[string][]*structs.NodeEvent{
		node.ID: {nodeEvent},
	}
	err = state.UpsertNodeEvents(structs.MsgTypeTestSetup, uint64(1001), nodeEvents)
	require.Nil(err)

	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.Nil(err)

	require.Equal(2, len(out.Events))
	require.Equal(nodeEvent, out.Events[1])
}

// To prevent stale node events from accumulating, we limit the number of
// stored node events to 10.
func TestStateStore_NodeEvents_RetentionWindow(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	node := mock.Node()

	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.Equal(1, len(node.Events))
	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)

	var out *structs.Node
	for i := 1; i <= 20; i++ {
		ws := memdb.NewWatchSet()
		out, err = state.NodeByID(ws, node.ID)
		require.Nil(err)

		nodeEvent := &structs.NodeEvent{
			Message:   fmt.Sprintf("%dith failed", i),
			Subsystem: "Driver",
			Timestamp: time.Now(),
		}

		nodeEvents := map[string][]*structs.NodeEvent{
			out.ID: {nodeEvent},
		}
		err := state.UpsertNodeEvents(structs.MsgTypeTestSetup, uint64(i), nodeEvents)
		require.Nil(err)

		require.True(watchFired(ws))
		ws = memdb.NewWatchSet()
		out, err = state.NodeByID(ws, node.ID)
		require.Nil(err)
	}

	ws := memdb.NewWatchSet()
	out, err = state.NodeByID(ws, node.ID)
	require.Nil(err)

	require.Equal(10, len(out.Events))
	require.Equal(uint64(11), out.Events[0].CreateIndex)
	require.Equal(uint64(20), out.Events[len(out.Events)-1].CreateIndex)
}

func TestStateStore_UpdateNodeDrain_ResetEligiblity(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	node := mock.Node()
	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))

	// Create a watchset so we can test that update node drain fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.NodeByID(ws, node.ID)
	require.Nil(err)

	drain := &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: -1 * time.Second,
		},
	}

	event1 := &structs.NodeEvent{
		Message:   "Drain strategy enabled",
		Subsystem: structs.NodeEventSubsystemDrain,
		Timestamp: time.Now(),
	}
	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1001, node.ID, drain, false, 7, event1))
	require.True(watchFired(ws))

	// Remove the drain
	event2 := &structs.NodeEvent{
		Message:   "Drain strategy disabled",
		Subsystem: structs.NodeEventSubsystemDrain,
		Timestamp: time.Now(),
	}
	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1002, node.ID, nil, true, 9, event2))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.Nil(err)
	require.False(out.Drain)
	require.Nil(out.DrainStrategy)
	require.Equal(out.SchedulingEligibility, structs.NodeSchedulingEligible)
	require.Len(out.Events, 3)
	require.EqualValues(1002, out.ModifyIndex)
	require.EqualValues(9, out.StatusUpdatedAt)

	index, err := state.Index("nodes")
	require.Nil(err)
	require.EqualValues(1002, index)
	require.False(watchFired(ws))
}

func TestStateStore_UpdateNodeEligibility(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	node := mock.Node()

	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	expectedEligibility := structs.NodeSchedulingIneligible

	// Create a watchset so we can test that update node drain fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.NodeByID(ws, node.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	event := &structs.NodeEvent{
		Message:   "Node marked as ineligible",
		Subsystem: structs.NodeEventSubsystemCluster,
		Timestamp: time.Now(),
	}
	require.Nil(state.UpdateNodeEligibility(structs.MsgTypeTestSetup, 1001, node.ID, expectedEligibility, 7, event))
	require.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	require.Nil(err)
	require.Equal(out.SchedulingEligibility, expectedEligibility)
	require.Len(out.Events, 2)
	require.Equal(out.Events[1], event)
	require.EqualValues(1001, out.ModifyIndex)
	require.EqualValues(7, out.StatusUpdatedAt)

	index, err := state.Index("nodes")
	require.Nil(err)
	require.EqualValues(1001, index)
	require.False(watchFired(ws))

	// Set a drain strategy
	expectedDrain := &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: -1 * time.Second,
		},
	}
	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1002, node.ID, expectedDrain, false, 7, nil))

	// Try to set the node to eligible
	err = state.UpdateNodeEligibility(structs.MsgTypeTestSetup, 1003, node.ID, structs.NodeSchedulingEligible, 9, nil)
	require.NotNil(err)
	require.Contains(err.Error(), "while it is draining")
}

func TestStateStore_Nodes(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var nodes []*structs.Node

	for i := 0; i < 10; i++ {
		node := mock.Node()
		nodes = append(nodes, node)

		err := state.UpsertNode(structs.MsgTypeTestSetup, 1000+uint64(i), node)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	// Create a watchset so we can test that getters don't cause it to fire
	ws := memdb.NewWatchSet()
	iter, err := state.Nodes(ws)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	var out []*structs.Node
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Node))
	}

	sort.Sort(NodeIDSort(nodes))
	sort.Sort(NodeIDSort(out))

	if !reflect.DeepEqual(nodes, out) {
		t.Fatalf("bad: %#v %#v", nodes, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_NodesByIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	node := mock.Node()

	node.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that getters don't cause it to fire
	ws := memdb.NewWatchSet()
	iter, err := state.NodesByIDPrefix(ws, node.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	gatherNodes := func(iter memdb.ResultIterator) []*structs.Node {
		var nodes []*structs.Node
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			node := raw.(*structs.Node)
			nodes = append(nodes, node)
		}
		return nodes
	}

	nodes := gatherNodes(iter)
	if len(nodes) != 1 {
		t.Fatalf("err: %v", err)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}

	iter, err = state.NodesByIDPrefix(ws, "11")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	nodes = gatherNodes(iter)
	if len(nodes) != 1 {
		t.Fatalf("err: %v", err)
	}

	node = mock.Node()
	node.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
	err = state.UpsertNode(structs.MsgTypeTestSetup, 1001, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	iter, err = state.NodesByIDPrefix(ws, "11")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	nodes = gatherNodes(iter)
	if len(nodes) != 2 {
		t.Fatalf("err: %v", err)
	}

	iter, err = state.NodesByIDPrefix(ws, "1111")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	nodes = gatherNodes(iter)
	if len(nodes) != 1 {
		t.Fatalf("err: %v", err)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreNode(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	node := mock.Node()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.NodeRestore(node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.NodeByID(ws, node.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, node) {
		t.Fatalf("Bad: %#v %#v", out, node)
	}
}

func TestStateStore_UpsertJob_Job(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
		t.Fatalf("err: %v", err)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(job, out) {
		t.Fatalf("bad: %#v %#v", job, out)
	}

	index, err := state.Index("jobs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != job.ID {
		t.Fatalf("bad summary id: %v", summary.JobID)
	}
	_, ok := summary.Summary["web"]
	if !ok {
		t.Fatalf("nil summary for task group")
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}

	// Check the job versions
	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(allVersions) != 1 {
		t.Fatalf("got %d; want 1", len(allVersions))
	}

	if a := allVersions[0]; a.ID != job.ID || a.Version != 0 {
		t.Fatalf("bad: %v", a)
	}

	// Test the looking up the job by version returns the same results
	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, vout) {
		t.Fatalf("bad: %#v %#v", out, vout)
	}
}

func TestStateStore_UpdateUpsertJob_Job(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	job2 := mock.Job()
	job2.ID = job.ID
	job2.AllAtOnce = true
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(job2, out) {
		t.Fatalf("bad: %#v %#v", job2, out)
	}

	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1001 {
		t.Fatalf("bad: %#v", out)
	}
	if out.Version != 1 {
		t.Fatalf("bad: %#v", out)
	}

	index, err := state.Index("jobs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	// Test the looking up the job by version returns the same results
	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 1)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, vout) {
		t.Fatalf("bad: %#v %#v", out, vout)
	}

	// Test that the job summary remains the same if the job is updated but
	// count remains same
	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != job.ID {
		t.Fatalf("bad summary id: %v", summary.JobID)
	}
	_, ok := summary.Summary["web"]
	if !ok {
		t.Fatalf("nil summary for task group")
	}

	// Check the job versions
	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(allVersions) != 2 {
		t.Fatalf("got %d; want 1", len(allVersions))
	}

	if a := allVersions[0]; a.ID != job.ID || a.Version != 1 || !a.AllAtOnce {
		t.Fatalf("bad: %+v", a)
	}
	if a := allVersions[1]; a.ID != job.ID || a.Version != 0 || a.AllAtOnce {
		t.Fatalf("bad: %+v", a)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpdateUpsertJob_PeriodicJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.PeriodicJob()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a child and an evaluation
	job2 := job.Copy()
	job2.Periodic = nil
	job2.ID = fmt.Sprintf("%v/%s-1490635020", job.ID, structs.PeriodicLaunchSuffix)
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	eval := mock.Eval()
	eval.JobID = job2.ID
	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1002, []*structs.Evaluation{eval})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	job3 := job.Copy()
	job3.TaskGroups[0].Tasks[0].Name = "new name"
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1003, job3)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if s, e := out.Status, structs.JobStatusRunning; s != e {
		t.Fatalf("got status %v; want %v", s, e)
	}

}

func TestStateStore_UpsertJob_BadNamespace(t *testing.T) {
	t.Parallel()

	assert := assert.New(t)
	state := testStateStore(t)
	job := mock.Job()
	job.Namespace = "foo"

	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	assert.Contains(err.Error(), "nonexistent namespace")

	ws := memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	assert.Nil(err)
	assert.Nil(out)
}

// Upsert a job that is the child of a parent job and ensures its summary gets
// updated.
func TestStateStore_UpsertJob_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create a watchset so we can test that upsert fires the watch
	parent := mock.Job()
	ws := memdb.NewWatchSet()
	_, err := state.JobByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, parent); err != nil {
		t.Fatalf("err: %v", err)
	}

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1001, child); err != nil {
		t.Fatalf("err: %v", err)
	}

	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 1 || summary.Children.Running != 0 || summary.Children.Dead != 0 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpdateUpsertJob_JobVersion(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create a job and mark it as stable
	job := mock.Job()
	job.Stable = true
	job.Name = "0"

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	var finalJob *structs.Job
	for i := 1; i < 300; i++ {
		finalJob = mock.Job()
		finalJob.ID = job.ID
		finalJob.Name = fmt.Sprintf("%d", i)
		err = state.UpsertJob(structs.MsgTypeTestSetup, uint64(1000+i), finalJob)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(finalJob, out) {
		t.Fatalf("bad: %#v %#v", finalJob, out)
	}

	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1299 {
		t.Fatalf("bad: %#v", out)
	}
	if out.Version != 299 {
		t.Fatalf("bad: %#v", out)
	}

	index, err := state.Index("job_version")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1299 {
		t.Fatalf("bad: %d", index)
	}

	// Check the job versions
	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(allVersions) != structs.JobTrackedVersions {
		t.Fatalf("got %d; want %d", len(allVersions), structs.JobTrackedVersions)
	}

	if a := allVersions[0]; a.ID != job.ID || a.Version != 299 || a.Name != "299" {
		t.Fatalf("bad: %+v", a)
	}
	if a := allVersions[1]; a.ID != job.ID || a.Version != 298 || a.Name != "298" {
		t.Fatalf("bad: %+v", a)
	}

	// Ensure we didn't delete the stable job
	if a := allVersions[structs.JobTrackedVersions-1]; a.ID != job.ID ||
		a.Version != 0 || a.Name != "0" || !a.Stable {
		t.Fatalf("bad: %+v", a)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteJob_Job(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.JobByID(ws, job.Namespace, job.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.DeleteJob(1001, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", job, out)
	}

	index, err := state.Index("jobs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary != nil {
		t.Fatalf("expected summary to be nil, but got: %v", summary)
	}

	index, err = state.Index("job_summary")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if len(versions) != 0 {
		t.Fatalf("expected no job versions")
	}

	index, err = state.Index("job_summary")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteJobTxn_BatchDeletes(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	const testJobCount = 10
	const jobVersionCount = 4

	stateIndex := uint64(1000)

	jobs := make([]*structs.Job, testJobCount)
	for i := 0; i < testJobCount; i++ {
		stateIndex++
		job := mock.BatchJob()

		err := state.UpsertJob(structs.MsgTypeTestSetup, stateIndex, job)
		require.NoError(t, err)

		jobs[i] = job

		// Create some versions
		for vi := 1; vi < jobVersionCount; vi++ {
			stateIndex++

			job := job.Copy()
			job.TaskGroups[0].Tasks[0].Env = map[string]string{
				"Version": fmt.Sprintf("%d", vi),
			}

			require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, stateIndex, job))
		}
	}

	ws := memdb.NewWatchSet()

	// Sanity check that jobs are present in DB
	job, err := state.JobByID(ws, jobs[0].Namespace, jobs[0].ID)
	require.NoError(t, err)
	require.Equal(t, jobs[0].ID, job.ID)

	jobVersions, err := state.JobVersionsByID(ws, jobs[0].Namespace, jobs[0].ID)
	require.NoError(t, err)
	require.Equal(t, jobVersionCount, len(jobVersions))

	// Actually delete
	const deletionIndex = uint64(10001)
	err = state.WithWriteTransaction(structs.MsgTypeTestSetup, deletionIndex, func(txn Txn) error {
		for i, job := range jobs {
			err := state.DeleteJobTxn(deletionIndex, job.Namespace, job.ID, txn)
			require.NoError(t, err, "failed at %d %e", i, err)
		}
		return nil
	})
	assert.NoError(t, err)

	assert.True(t, watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, jobs[0].Namespace, jobs[0].ID)
	require.NoError(t, err)
	require.Nil(t, out)

	jobVersions, err = state.JobVersionsByID(ws, jobs[0].Namespace, jobs[0].ID)
	require.NoError(t, err)
	require.Empty(t, jobVersions)

	index, err := state.Index("jobs")
	require.NoError(t, err)
	require.Equal(t, deletionIndex, index)
}

func TestStateStore_DeleteJob_MultipleVersions(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	assert := assert.New(t)

	// Create a job and mark it as stable
	job := mock.Job()
	job.Stable = true
	job.Priority = 0

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	assert.Nil(err)
	assert.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1000, job))
	assert.True(watchFired(ws))

	var finalJob *structs.Job
	for i := 1; i < 20; i++ {
		finalJob = mock.Job()
		finalJob.ID = job.ID
		finalJob.Priority = i
		assert.Nil(state.UpsertJob(structs.MsgTypeTestSetup, uint64(1000+i), finalJob))
	}

	assert.Nil(state.DeleteJob(1020, job.Namespace, job.ID))
	assert.True(watchFired(ws))

	ws = memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	assert.Nil(err)
	assert.Nil(out)

	index, err := state.Index("jobs")
	assert.Nil(err)
	assert.EqualValues(1020, index)

	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
	assert.Nil(err)
	assert.Nil(summary)

	index, err = state.Index("job_version")
	assert.Nil(err)
	assert.EqualValues(1020, index)

	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
	assert.Nil(err)
	assert.Len(versions, 0)

	index, err = state.Index("job_summary")
	assert.Nil(err)
	assert.EqualValues(1020, index)

	assert.False(watchFired(ws))
}

func TestStateStore_DeleteJob_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	parent := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
		t.Fatalf("err: %v", err)
	}

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err := state.DeleteJob(1001, child.Namespace, child.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Jobs(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var jobs []*structs.Job

	for i := 0; i < 10; i++ {
		job := mock.Job()
		jobs = append(jobs, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.Jobs(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Job))
	}

	sort.Sort(JobIDSort(jobs))
	sort.Sort(JobIDSort(out))

	if !reflect.DeepEqual(jobs, out) {
		t.Fatalf("bad: %#v %#v", jobs, out)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_JobVersions(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var jobs []*structs.Job

	for i := 0; i < 10; i++ {
		job := mock.Job()
		jobs = append(jobs, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.JobVersions(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Job))
	}

	sort.Sort(JobIDSort(jobs))
	sort.Sort(JobIDSort(out))

	if !reflect.DeepEqual(jobs, out) {
		t.Fatalf("bad: %#v %#v", jobs, out)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_JobsByIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	job.ID = "redis"
	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	iter, err := state.JobsByIDPrefix(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	gatherJobs := func(iter memdb.ResultIterator) []*structs.Job {
		var jobs []*structs.Job
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			jobs = append(jobs, raw.(*structs.Job))
		}
		return jobs
	}

	jobs := gatherJobs(iter)
	if len(jobs) != 1 {
		t.Fatalf("err: %v", err)
	}

	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "re")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	jobs = gatherJobs(iter)
	if len(jobs) != 1 {
		t.Fatalf("err: %v", err)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}

	job = mock.Job()
	job.ID = "riak"
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "r")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	jobs = gatherJobs(iter)
	if len(jobs) != 2 {
		t.Fatalf("err: %v", err)
	}

	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "ri")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	jobs = gatherJobs(iter)
	if len(jobs) != 1 {
		t.Fatalf("err: %v", err)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_JobsByPeriodic(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var periodic, nonPeriodic []*structs.Job

	for i := 0; i < 10; i++ {
		job := mock.Job()
		nonPeriodic = append(nonPeriodic, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	for i := 0; i < 10; i++ {
		job := mock.PeriodicJob()
		periodic = append(periodic, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.JobsByPeriodic(ws, true)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var outPeriodic []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		outPeriodic = append(outPeriodic, raw.(*structs.Job))
	}

	iter, err = state.JobsByPeriodic(ws, false)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var outNonPeriodic []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		outNonPeriodic = append(outNonPeriodic, raw.(*structs.Job))
	}

	sort.Sort(JobIDSort(periodic))
	sort.Sort(JobIDSort(nonPeriodic))
	sort.Sort(JobIDSort(outPeriodic))
	sort.Sort(JobIDSort(outNonPeriodic))

	if !reflect.DeepEqual(periodic, outPeriodic) {
		t.Fatalf("bad: %#v %#v", periodic, outPeriodic)
	}

	if !reflect.DeepEqual(nonPeriodic, outNonPeriodic) {
		t.Fatalf("bad: %#v %#v", nonPeriodic, outNonPeriodic)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_JobsByScheduler(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var serviceJobs []*structs.Job
	var sysJobs []*structs.Job

	for i := 0; i < 10; i++ {
		job := mock.Job()
		serviceJobs = append(serviceJobs, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	for i := 0; i < 10; i++ {
		job := mock.SystemJob()
		job.Status = structs.JobStatusRunning
		sysJobs = append(sysJobs, job)

		err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.JobsByScheduler(ws, "service")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var outService []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		outService = append(outService, raw.(*structs.Job))
	}

	iter, err = state.JobsByScheduler(ws, "system")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var outSystem []*structs.Job
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		outSystem = append(outSystem, raw.(*structs.Job))
	}

	sort.Sort(JobIDSort(serviceJobs))
	sort.Sort(JobIDSort(sysJobs))
	sort.Sort(JobIDSort(outService))
	sort.Sort(JobIDSort(outSystem))

	if !reflect.DeepEqual(serviceJobs, outService) {
		t.Fatalf("bad: %#v %#v", serviceJobs, outService)
	}

	if !reflect.DeepEqual(sysJobs, outSystem) {
		t.Fatalf("bad: %#v %#v", sysJobs, outSystem)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_JobsByGC(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	gc, nonGc := make(map[string]struct{}), make(map[string]struct{})

	for i := 0; i < 20; i++ {
		var job *structs.Job
		if i%2 == 0 {
			job = mock.Job()
		} else {
			job = mock.PeriodicJob()
		}
		nonGc[job.ID] = struct{}{}

		if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job); err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	for i := 0; i < 20; i += 2 {
		job := mock.Job()
		job.Type = structs.JobTypeBatch
		gc[job.ID] = struct{}{}

		if err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job); err != nil {
			t.Fatalf("err: %v", err)
		}

		// Create an eval for it
		eval := mock.Eval()
		eval.JobID = job.ID
		eval.Status = structs.EvalStatusComplete
		if err := state.UpsertEvals(structs.MsgTypeTestSetup, 2000+uint64(i+1), []*structs.Evaluation{eval}); err != nil {
			t.Fatalf("err: %v", err)
		}

	}

	ws := memdb.NewWatchSet()
	iter, err := state.JobsByGC(ws, true)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	outGc := make(map[string]struct{})
	for i := iter.Next(); i != nil; i = iter.Next() {
		j := i.(*structs.Job)
		outGc[j.ID] = struct{}{}
	}

	iter, err = state.JobsByGC(ws, false)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	outNonGc := make(map[string]struct{})
	for i := iter.Next(); i != nil; i = iter.Next() {
		j := i.(*structs.Job)
		outNonGc[j.ID] = struct{}{}
	}

	if !reflect.DeepEqual(gc, outGc) {
		t.Fatalf("bad: %#v %#v", gc, outGc)
	}

	if !reflect.DeepEqual(nonGc, outNonGc) {
		t.Fatalf("bad: %#v %#v", nonGc, outNonGc)
	}
	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.JobRestore(job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, job) {
		t.Fatalf("Bad: %#v %#v", out, job)
	}
}

func TestStateStore_UpsertPeriodicLaunch(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	launch := &structs.PeriodicLaunch{
		ID:        job.ID,
		Namespace: job.Namespace,
		Launch:    time.Now(),
	}

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err := state.UpsertPeriodicLaunch(1000, launch)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}

	if !reflect.DeepEqual(launch, out) {
		t.Fatalf("bad: %#v %#v", job, out)
	}

	index, err := state.Index("periodic_launch")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpdateUpsertPeriodicLaunch(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	launch := &structs.PeriodicLaunch{
		ID:        job.ID,
		Namespace: job.Namespace,
		Launch:    time.Now(),
	}

	err := state.UpsertPeriodicLaunch(1000, launch)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	launch2 := &structs.PeriodicLaunch{
		ID:        job.ID,
		Namespace: job.Namespace,
		Launch:    launch.Launch.Add(1 * time.Second),
	}
	err = state.UpsertPeriodicLaunch(1001, launch2)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1001 {
		t.Fatalf("bad: %#v", out)
	}

	if !reflect.DeepEqual(launch2, out) {
		t.Fatalf("bad: %#v %#v", launch2, out)
	}

	index, err := state.Index("periodic_launch")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeletePeriodicLaunch(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	launch := &structs.PeriodicLaunch{
		ID:        job.ID,
		Namespace: job.Namespace,
		Launch:    time.Now(),
	}

	err := state.UpsertPeriodicLaunch(1000, launch)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.DeletePeriodicLaunch(1001, launch.Namespace, launch.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", job, out)
	}

	index, err := state.Index("periodic_launch")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_PeriodicLaunches(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var launches []*structs.PeriodicLaunch

	for i := 0; i < 10; i++ {
		job := mock.Job()
		launch := &structs.PeriodicLaunch{
			ID:        job.ID,
			Namespace: job.Namespace,
			Launch:    time.Now(),
		}
		launches = append(launches, launch)

		err := state.UpsertPeriodicLaunch(1000+uint64(i), launch)
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.PeriodicLaunches(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	out := make(map[string]*structs.PeriodicLaunch, 10)
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		launch := raw.(*structs.PeriodicLaunch)
		if _, ok := out[launch.ID]; ok {
			t.Fatalf("duplicate: %v", launch.ID)
		}

		out[launch.ID] = launch
	}

	for _, launch := range launches {
		l, ok := out[launch.ID]
		if !ok {
			t.Fatalf("bad %v", launch.ID)
		}

		if !reflect.DeepEqual(launch, l) {
			t.Fatalf("bad: %#v %#v", launch, l)
		}

		delete(out, launch.ID)
	}

	if len(out) != 0 {
		t.Fatalf("leftover: %#v", out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestorePeriodicLaunch(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	launch := &structs.PeriodicLaunch{
		ID:        job.ID,
		Namespace: job.Namespace,
		Launch:    time.Now(),
	}

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.PeriodicLaunchRestore(launch)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, launch) {
		t.Fatalf("Bad: %#v %#v", out, job)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreJobVersion(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.JobVersionRestore(job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, job.Version)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, job) {
		t.Fatalf("Bad: %#v %#v", out, job)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreDeployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	d := mock.Deployment()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.DeploymentRestore(d)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.DeploymentByID(ws, d.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, d) {
		t.Fatalf("Bad: %#v %#v", out, d)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreJobSummary(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	jobSummary := &structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Starting: 10,
			},
		},
	}
	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.JobSummaryRestore(jobSummary)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, jobSummary) {
		t.Fatalf("Bad: %#v %#v", out, jobSummary)
	}
}

// TestStateStore_CSIVolume checks register, list and deregister for csi_volumes
func TestStateStore_CSIVolume(t *testing.T) {
	state := testStateStore(t)
	index := uint64(1000)

	// Volume IDs
	vol0, vol1 := uuid.Generate(), uuid.Generate()

	// Create a node running a healthy instance of the plugin
	node := mock.Node()
	pluginID := "minnie"
	alloc := mock.Alloc()
	alloc.DesiredStatus = "run"
	alloc.ClientStatus = "running"
	alloc.NodeID = node.ID
	alloc.Job.TaskGroups[0].Volumes = map[string]*structs.VolumeRequest{
		"foo": {
			Name:   "foo",
			Source: vol0,
			Type:   "csi",
		},
	}

	node.CSINodePlugins = map[string]*structs.CSIInfo{
		pluginID: {
			PluginID:                 pluginID,
			AllocID:                  alloc.ID,
			Healthy:                  true,
			HealthDescription:        "healthy",
			RequiresControllerPlugin: false,
			RequiresTopologies:       false,
			NodeInfo: &structs.CSINodeInfo{
				ID:                      node.ID,
				MaxVolumes:              64,
				RequiresNodeStageVolume: true,
			},
		},
	}

	index++
	err := state.UpsertNode(structs.MsgTypeTestSetup, index, node)
	require.NoError(t, err)
	defer state.DeleteNode(structs.MsgTypeTestSetup, 9999, []string{pluginID})

	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc})
	require.NoError(t, err)

	ns := structs.DefaultNamespace

	v0 := structs.NewCSIVolume("foo", index)
	v0.ID = vol0
	v0.Namespace = ns
	v0.PluginID = "minnie"
	v0.Schedulable = true
	v0.AccessMode = structs.CSIVolumeAccessModeMultiNodeSingleWriter
	v0.AttachmentMode = structs.CSIVolumeAttachmentModeFilesystem

	index++
	v1 := structs.NewCSIVolume("foo", index)
	v1.ID = vol1
	v1.Namespace = ns
	v1.PluginID = "adam"
	v1.Schedulable = true
	v1.AccessMode = structs.CSIVolumeAccessModeMultiNodeSingleWriter
	v1.AttachmentMode = structs.CSIVolumeAttachmentModeFilesystem

	index++
	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0, v1})
	require.NoError(t, err)

	// volume registration is idempotent, unless identies are changed
	index++
	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0, v1})
	require.NoError(t, err)

	index++
	v2 := v0.Copy()
	v2.PluginID = "new-id"
	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v2})
	require.Error(t, err, fmt.Sprintf("volume exists: %s", v0.ID))

	ws := memdb.NewWatchSet()
	iter, err := state.CSIVolumesByNamespace(ws, ns)
	require.NoError(t, err)

	slurp := func(iter memdb.ResultIterator) (vs []*structs.CSIVolume) {
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			vol := raw.(*structs.CSIVolume)
			vs = append(vs, vol)
		}
		return vs
	}

	vs := slurp(iter)
	require.Equal(t, 2, len(vs))

	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
	require.NoError(t, err)
	vs = slurp(iter)
	require.Equal(t, 1, len(vs))

	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByNodeID(ws, node.ID)
	require.NoError(t, err)
	vs = slurp(iter)
	require.Equal(t, 1, len(vs))

	// Allocs
	a0 := mock.Alloc()
	a1 := mock.Alloc()
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{a0, a1})
	require.NoError(t, err)

	// Claims
	r := structs.CSIVolumeClaimRead
	w := structs.CSIVolumeClaimWrite
	u := structs.CSIVolumeClaimGC
	claim0 := &structs.CSIVolumeClaim{
		AllocationID: a0.ID,
		NodeID:       node.ID,
		Mode:         r,
	}
	claim1 := &structs.CSIVolumeClaim{
		AllocationID: a1.ID,
		NodeID:       node.ID,
		Mode:         w,
	}

	index++
	err = state.CSIVolumeClaim(index, ns, vol0, claim0)
	require.NoError(t, err)
	index++
	err = state.CSIVolumeClaim(index, ns, vol0, claim1)
	require.NoError(t, err)

	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
	require.NoError(t, err)
	vs = slurp(iter)
	require.False(t, vs[0].WriteFreeClaims())

	claim0.Mode = u
	err = state.CSIVolumeClaim(2, ns, vol0, claim0)
	require.NoError(t, err)
	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
	require.NoError(t, err)
	vs = slurp(iter)
	require.True(t, vs[0].ReadSchedulable())

	// registration is an error when the volume is in use
	index++
	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0})
	require.Error(t, err, fmt.Sprintf("volume exists: %s", vol0))
	// as is deregistration
	index++
	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, false)
	require.Error(t, err, fmt.Sprintf("volume in use: %s", vol0))

	// even if forced, because we have a non-terminal claim
	index++
	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, true)
	require.Error(t, err, fmt.Sprintf("volume in use: %s", vol0))

	// release claims to unblock deregister
	index++
	claim0.State = structs.CSIVolumeClaimStateReadyToFree
	err = state.CSIVolumeClaim(index, ns, vol0, claim0)
	require.NoError(t, err)
	index++
	claim1.Mode = u
	claim1.State = structs.CSIVolumeClaimStateReadyToFree
	err = state.CSIVolumeClaim(index, ns, vol0, claim1)
	require.NoError(t, err)

	index++
	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, false)
	require.NoError(t, err)

	// List, now omitting the deregistered volume
	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
	require.NoError(t, err)
	vs = slurp(iter)
	require.Equal(t, 0, len(vs))

	ws = memdb.NewWatchSet()
	iter, err = state.CSIVolumesByNamespace(ws, ns)
	require.NoError(t, err)
	vs = slurp(iter)
	require.Equal(t, 1, len(vs))
}

// TestStateStore_CSIPluginNodes uses node fingerprinting to create a plugin and update health
func TestStateStore_CSIPluginNodes(t *testing.T) {
	index := uint64(999)
	state := testStateStore(t)
	ws := memdb.NewWatchSet()
	plugID := "foo"

	// Create Nomad client Nodes
	ns := []*structs.Node{mock.Node(), mock.Node()}
	for _, n := range ns {
		index++
		err := state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)
	}

	// Fingerprint a running controller plugin
	n0, _ := state.NodeByID(ws, ns[0].ID)
	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			Healthy:                  true,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			ControllerInfo: &structs.CSIControllerInfo{
				SupportsReadOnlyAttach: true,
				SupportsListVolumes:    true,
			},
		},
	}
	index++
	err := state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	// Fingerprint two running node plugins
	for _, n := range ns[:] {
		n, _ := state.NodeByID(ws, n.ID)
		n.CSINodePlugins = map[string]*structs.CSIInfo{
			plugID: {
				PluginID:                 plugID,
				Healthy:                  true,
				UpdateTime:               time.Now(),
				RequiresControllerPlugin: true,
				RequiresTopologies:       false,
				NodeInfo:                 &structs.CSINodeInfo{},
			},
		}
		index++
		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)
	}

	plug, err := state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.True(t, plug.ControllerRequired)
	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

	// Volume using the plugin
	index++
	vol := &structs.CSIVolume{
		ID:        uuid.Generate(),
		Namespace: structs.DefaultNamespace,
		PluginID:  plugID,
	}
	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{vol})
	require.NoError(t, err)

	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
	require.NoError(t, err)
	require.True(t, vol.Schedulable, "volume should be schedulable")

	// Controller is unhealthy
	n0, _ = state.NodeByID(ws, ns[0].ID)
	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			Healthy:                  false,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			ControllerInfo: &structs.CSIControllerInfo{
				SupportsReadOnlyAttach: true,
				SupportsListVolumes:    true,
			},
		},
	}

	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
	require.NoError(t, err)
	require.False(t, vol.Schedulable, "volume should not be schedulable")

	// Node plugin is removed
	n1, _ := state.NodeByID(ws, ns[1].ID)
	n1.CSINodePlugins = map[string]*structs.CSIInfo{}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n1)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 1, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 1, len(plug.Nodes), "nodes expected")

	// Last node plugin is removed
	n0, _ = state.NodeByID(ws, ns[0].ID)
	n0.CSINodePlugins = map[string]*structs.CSIInfo{}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	// Nodes plugins should be gone but controllers left
	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 0, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 0, len(plug.Nodes), "nodes expected")

	// A node plugin is restored
	n0, _ = state.NodeByID(ws, n0.ID)
	n0.CSINodePlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			Healthy:                  true,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			NodeInfo:                 &structs.CSINodeInfo{},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	// Nodes plugin should be replaced and healthy
	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 1, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 1, len(plug.Nodes), "nodes expected")

	// Remove node again
	n0, _ = state.NodeByID(ws, ns[0].ID)
	n0.CSINodePlugins = map[string]*structs.CSIInfo{}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	// Nodes plugins should be gone but controllers left
	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 0, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 0, len(plug.Nodes), "nodes expected")

	// controller is removed
	n0, _ = state.NodeByID(ws, ns[0].ID)
	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	// Plugin has been removed entirely
	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Nil(t, plug)

	// Volume still exists and is safe to query, but unschedulable
	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
	require.NoError(t, err)
	require.False(t, vol.Schedulable)
}

// TestStateStore_CSIPluginAllocUpdates tests the ordering
// interactions for CSI plugins between Nomad client node updates and
// allocation updates.
func TestStateStore_CSIPluginAllocUpdates(t *testing.T) {
	t.Parallel()
	index := uint64(999)
	state := testStateStore(t)
	ws := memdb.NewWatchSet()

	n := mock.Node()
	index++
	err := state.UpsertNode(structs.MsgTypeTestSetup, index, n)
	require.NoError(t, err)

	// (1) unhealthy fingerprint, then terminal alloc, then healthy node update
	plugID0 := "foo0"

	alloc0 := mock.Alloc()
	alloc0.NodeID = n.ID
	alloc0.DesiredStatus = "run"
	alloc0.ClientStatus = "running"
	alloc0.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID0}
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
	require.NoError(t, err)

	n, _ = state.NodeByID(ws, n.ID)
	n.CSINodePlugins = map[string]*structs.CSIInfo{
		plugID0: {
			PluginID:                 plugID0,
			AllocID:                  alloc0.ID,
			Healthy:                  false,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			NodeInfo:                 &structs.CSINodeInfo{},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
	require.NoError(t, err)

	plug, err := state.CSIPluginByID(ws, plugID0)
	require.NoError(t, err)
	require.Nil(t, plug, "no plugin should exist: not yet healthy")

	alloc0.DesiredStatus = "stopped"
	alloc0.ClientStatus = "complete"
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID0)
	require.NoError(t, err)
	require.Nil(t, plug, "no plugin should exist: allocs never healthy")

	n, _ = state.NodeByID(ws, n.ID)
	n.CSINodePlugins[plugID0].Healthy = true
	n.CSINodePlugins[plugID0].UpdateTime = time.Now()
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID0)
	require.NoError(t, err)
	require.NotNil(t, plug, "plugin should exist")

	// (2) healthy fingerprint, then terminal alloc update
	plugID1 := "foo1"

	alloc1 := mock.Alloc()
	n, _ = state.NodeByID(ws, n.ID)
	n.CSINodePlugins = map[string]*structs.CSIInfo{
		plugID1: {
			PluginID:                 plugID1,
			AllocID:                  alloc1.ID,
			Healthy:                  true,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			NodeInfo:                 &structs.CSINodeInfo{},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID1)
	require.NoError(t, err)
	require.NotNil(t, plug, "plugin should exist")

	alloc1.NodeID = n.ID
	alloc1.DesiredStatus = "stop"
	alloc1.ClientStatus = "complete"
	alloc1.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID1}
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID1)
	require.NoError(t, err)
	require.Nil(t, plug, "no plugin should exist: alloc became terminal")

	// (3) terminal alloc update, then unhealthy fingerprint
	plugID2 := "foo2"

	alloc2 := mock.Alloc()
	alloc2.NodeID = n.ID
	alloc2.DesiredStatus = "stop"
	alloc2.ClientStatus = "complete"
	alloc2.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID2}
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc2})
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID2)
	require.NoError(t, err)
	require.Nil(t, plug, "no plugin should exist: alloc became terminal")

	n, _ = state.NodeByID(ws, n.ID)
	n.CSINodePlugins = map[string]*structs.CSIInfo{
		plugID2: {
			PluginID:                 plugID2,
			AllocID:                  alloc2.ID,
			Healthy:                  false,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			NodeInfo:                 &structs.CSINodeInfo{},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID2)
	require.NoError(t, err)
	require.Nil(t, plug, "plugin should not exist: never became healthy")

}

// TestStateStore_CSIPluginMultiNodeUpdates tests the ordering
// interactions for CSI plugins between Nomad client node updates and
// allocation updates when multiple nodes are involved
func TestStateStore_CSIPluginMultiNodeUpdates(t *testing.T) {
	t.Parallel()
	index := uint64(999)
	state := testStateStore(t)
	ws := memdb.NewWatchSet()

	var err error

	// Create Nomad client Nodes
	ns := []*structs.Node{mock.Node(), mock.Node()}
	for _, n := range ns {
		index++
		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)
	}

	plugID := "foo"
	plugCfg := &structs.TaskCSIPluginConfig{ID: plugID}

	// Fingerprint two running node plugins and their allocs; we'll
	// leave these in place for the test to ensure we don't GC the
	// plugin
	for _, n := range ns[:] {
		nAlloc := mock.Alloc()
		n, _ := state.NodeByID(ws, n.ID)
		n.CSINodePlugins = map[string]*structs.CSIInfo{
			plugID: {
				PluginID:                 plugID,
				AllocID:                  nAlloc.ID,
				Healthy:                  true,
				UpdateTime:               time.Now(),
				RequiresControllerPlugin: true,
				RequiresTopologies:       false,
				NodeInfo:                 &structs.CSINodeInfo{},
			},
		}
		index++
		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)

		nAlloc.NodeID = n.ID
		nAlloc.DesiredStatus = "run"
		nAlloc.ClientStatus = "running"
		nAlloc.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg

		index++
		err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{nAlloc})
		require.NoError(t, err)
	}

	// Fingerprint a running controller plugin
	alloc0 := mock.Alloc()
	n0, _ := state.NodeByID(ws, ns[0].ID)
	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			AllocID:                  alloc0.ID,
			Healthy:                  true,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			ControllerInfo: &structs.CSIControllerInfo{
				SupportsReadOnlyAttach: true,
				SupportsListVolumes:    true,
			},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	plug, err := state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

	n1, _ := state.NodeByID(ws, ns[1].ID)

	alloc0.NodeID = n0.ID
	alloc0.DesiredStatus = "stop"
	alloc0.ClientStatus = "complete"
	alloc0.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg

	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 0, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

	alloc1 := mock.Alloc()
	alloc1.NodeID = n1.ID
	alloc1.DesiredStatus = "run"
	alloc1.ClientStatus = "running"
	alloc1.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg

	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 0, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

	n0, _ = state.NodeByID(ws, ns[0].ID)
	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			AllocID:                  alloc0.ID,
			Healthy:                  false,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			ControllerInfo: &structs.CSIControllerInfo{
				SupportsReadOnlyAttach: true,
				SupportsListVolumes:    true,
			},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
	require.NoError(t, err)

	n1.CSIControllerPlugins = map[string]*structs.CSIInfo{
		plugID: {
			PluginID:                 plugID,
			AllocID:                  alloc1.ID,
			Healthy:                  true,
			UpdateTime:               time.Now(),
			RequiresControllerPlugin: true,
			RequiresTopologies:       false,
			ControllerInfo: &structs.CSIControllerInfo{
				SupportsReadOnlyAttach: true,
				SupportsListVolumes:    true,
			},
		},
	}
	index++
	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n1)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.True(t, plug.ControllerRequired)
	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 2, len(plug.Nodes), "nodes expected")

}

// TestStateStore_CSIPlugin_ConcurrentStop tests that concurrent allocation
// updates don't cause the count to drift unexpectedly or cause allocation
// update errors.
func TestStateStore_CSIPlugin_ConcurrentStop(t *testing.T) {
	t.Parallel()
	index := uint64(999)
	state := testStateStore(t)
	ws := memdb.NewWatchSet()

	var err error

	// Create Nomad client Nodes
	ns := []*structs.Node{mock.Node(), mock.Node(), mock.Node()}
	for _, n := range ns {
		index++
		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)
	}

	plugID := "foo"
	plugCfg := &structs.TaskCSIPluginConfig{ID: plugID}

	allocs := []*structs.Allocation{}

	// Fingerprint 3 running node plugins and their allocs
	for _, n := range ns[:] {
		alloc := mock.Alloc()
		n, _ := state.NodeByID(ws, n.ID)
		n.CSINodePlugins = map[string]*structs.CSIInfo{
			plugID: {
				PluginID:                 plugID,
				AllocID:                  alloc.ID,
				Healthy:                  true,
				UpdateTime:               time.Now(),
				RequiresControllerPlugin: true,
				RequiresTopologies:       false,
				NodeInfo:                 &structs.CSINodeInfo{},
			},
		}
		index++
		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
		require.NoError(t, err)

		alloc.NodeID = n.ID
		alloc.DesiredStatus = "run"
		alloc.ClientStatus = "running"
		alloc.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg

		index++
		err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc})
		require.NoError(t, err)

		allocs = append(allocs, alloc)
	}

	plug, err := state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Equal(t, 3, plug.NodesHealthy, "nodes healthy")
	require.Equal(t, 3, len(plug.Nodes), "nodes expected")

	// stop all the allocs
	for _, alloc := range allocs {
		alloc.DesiredStatus = "stop"
		alloc.ClientStatus = "complete"
	}

	// this is somewhat artificial b/c we get alloc updates from multiple
	// nodes concurrently but not in a single RPC call. But this guarantees
	// we'll trigger any nested transaction setup bugs
	index++
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, allocs)
	require.NoError(t, err)

	plug, err = state.CSIPluginByID(ws, plugID)
	require.NoError(t, err)
	require.Nil(t, plug)
}

func TestStateStore_CSIPluginJobs(t *testing.T) {
	s := testStateStore(t)
	index := uint64(1001)

	controllerJob := mock.Job()
	controllerJob.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{
		ID:   "foo",
		Type: structs.CSIPluginTypeController,
	}

	nodeJob := mock.Job()
	nodeJob.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{
		ID:   "foo",
		Type: structs.CSIPluginTypeNode,
	}

	err := s.UpsertJob(structs.MsgTypeTestSetup, index, controllerJob)
	require.NoError(t, err)
	index++

	err = s.UpsertJob(structs.MsgTypeTestSetup, index, nodeJob)
	require.NoError(t, err)
	index++

	// Get the plugin, and make better fake allocations for it
	ws := memdb.NewWatchSet()
	plug, err := s.CSIPluginByID(ws, "foo")
	require.NoError(t, err)
	index++

	as := []*structs.Allocation{}
	for id, info := range plug.Controllers {
		as = append(as, &structs.Allocation{
			ID:        info.AllocID,
			Namespace: controllerJob.Namespace,
			JobID:     controllerJob.ID,
			Job:       controllerJob,
			TaskGroup: "web",
			EvalID:    uuid.Generate(),
			NodeID:    id,
		})
	}
	for id, info := range plug.Nodes {
		as = append(as, &structs.Allocation{
			ID:        info.AllocID,
			JobID:     nodeJob.ID,
			Namespace: nodeJob.Namespace,
			Job:       nodeJob,
			TaskGroup: "web",
			EvalID:    uuid.Generate(),
			NodeID:    id,
		})
	}

	err = s.UpsertAllocs(structs.MsgTypeTestSetup, index, as)
	require.NoError(t, err)
	index++

	// We use the summary to add
	err = s.ReconcileJobSummaries(index)
	require.NoError(t, err)
	index++

	// Delete a job
	err = s.DeleteJob(index, controllerJob.Namespace, controllerJob.ID)
	require.NoError(t, err)
	index++

	// plugin still exists
	plug, err = s.CSIPluginByID(ws, "foo")
	require.NoError(t, err)
	require.NotNil(t, plug)
	require.Equal(t, 0, len(plug.Controllers))

	// Delete a job
	err = s.DeleteJob(index, nodeJob.Namespace, nodeJob.ID)
	require.NoError(t, err)
	index++

	// plugin was collected
	plug, err = s.CSIPluginByID(ws, "foo")
	require.NoError(t, err)
	require.True(t, plug.IsEmpty())
}

func TestStateStore_RestoreCSIPlugin(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	plugin := mock.CSIPlugin()

	restore, err := state.Restore()
	require.NoError(err)

	err = restore.CSIPluginRestore(plugin)
	require.NoError(err)
	require.NoError(restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.CSIPluginByID(ws, plugin.ID)
	require.NoError(err)
	require.EqualValues(out, plugin)
}

func TestStateStore_RestoreCSIVolume(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	plugin := mock.CSIPlugin()
	volume := mock.CSIVolume(plugin)

	restore, err := state.Restore()
	require.NoError(err)

	err = restore.CSIVolumeRestore(volume)
	require.NoError(err)
	restore.Commit()

	ws := memdb.NewWatchSet()
	out, err := state.CSIVolumeByID(ws, "default", volume.ID)
	require.NoError(err)
	require.EqualValues(out, volume)
}

func TestStateStore_Indexes(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	node := mock.Node()

	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	iter, err := state.Indexes()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*IndexEntry
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*IndexEntry))
	}

	expect := &IndexEntry{"nodes", 1000}
	if l := len(out); l < 1 {
		t.Fatalf("unexpected number of index entries: %v", pretty.Sprint(out))
	}

	for _, index := range out {
		if index.Key != expect.Key {
			continue
		}
		if index.Value != expect.Value {
			t.Fatalf("bad index; got %d; want %d", index.Value, expect.Value)
		}

		// We matched
		return
	}

	t.Fatal("did not find expected index entry")
}

func TestStateStore_LatestIndex(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	if err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, mock.Node()); err != nil {
		t.Fatalf("err: %v", err)
	}

	exp := uint64(2000)
	if err := state.UpsertJob(structs.MsgTypeTestSetup, exp, mock.Job()); err != nil {
		t.Fatalf("err: %v", err)
	}

	latest, err := state.LatestIndex()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if latest != exp {
		t.Fatalf("LatestIndex() returned %d; want %d", latest, exp)
	}
}

func TestStateStore_RestoreIndex(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	index := &IndexEntry{"jobs", 1000}
	err = restore.IndexRestore(index)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	require.NoError(t, restore.Commit())

	out, err := state.Index("jobs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != 1000 {
		t.Fatalf("Bad: %#v %#v", out, 1000)
	}
}

func TestStateStore_UpsertEvals_Eval(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	eval := mock.Eval()

	// Create a watchset so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.EvalByID(ws, eval.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(eval, out) {
		t.Fatalf("bad: %#v %#v", eval, out)
	}

	index, err := state.Index("evals")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertEvals_CancelBlocked(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create two blocked evals for the same job
	j := "test-job"
	b1, b2 := mock.Eval(), mock.Eval()
	b1.JobID = j
	b1.Status = structs.EvalStatusBlocked
	b2.JobID = j
	b2.Status = structs.EvalStatusBlocked

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 999, []*structs.Evaluation{b1, b2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create one complete and successful eval for the job
	eval := mock.Eval()
	eval.JobID = j
	eval.Status = structs.EvalStatusComplete

	// Create a watchset so we can test that the upsert of the complete eval
	// fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.EvalByID(ws, b1.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalByID(ws, b2.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(eval, out) {
		t.Fatalf("bad: %#v %#v", eval, out)
	}

	index, err := state.Index("evals")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	// Get b1/b2 and check they are cancelled
	out1, err := state.EvalByID(ws, b1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	out2, err := state.EvalByID(ws, b2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out1.Status != structs.EvalStatusCancelled || out2.Status != structs.EvalStatusCancelled {
		t.Fatalf("bad: %#v %#v", out1, out2)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Update_UpsertEvals_Eval(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	eval := mock.Eval()

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watchset so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	ws2 := memdb.NewWatchSet()
	if _, err := state.EvalByID(ws, eval.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	if _, err := state.EvalsByJob(ws2, eval.Namespace, eval.JobID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	eval2 := mock.Eval()
	eval2.ID = eval.ID
	eval2.JobID = eval.JobID
	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}
	if !watchFired(ws2) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(eval2, out) {
		t.Fatalf("bad: %#v %#v", eval2, out)
	}

	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1001 {
		t.Fatalf("bad: %#v", out)
	}

	index, err := state.Index("evals")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertEvals_Eval_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	parent := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
		t.Fatalf("err: %v", err)
	}

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
		t.Fatalf("err: %v", err)
	}

	eval := mock.Eval()
	eval.Status = structs.EvalStatusComplete
	eval.JobID = child.ID

	// Create watchsets so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	ws2 := memdb.NewWatchSet()
	ws3 := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalByID(ws2, eval.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalsByJob(ws3, eval.Namespace, eval.JobID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}
	if !watchFired(ws2) {
		t.Fatalf("bad")
	}
	if !watchFired(ws3) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(eval, out) {
		t.Fatalf("bad: %#v %#v", eval, out)
	}

	index, err := state.Index("evals")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteEval_Eval(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	eval1 := mock.Eval()
	eval2 := mock.Eval()
	alloc1 := mock.Alloc()
	alloc2 := mock.Alloc()

	// Create watchsets so we can test that upsert fires the watch
	watches := make([]memdb.WatchSet, 12)
	for i := 0; i < 12; i++ {
		watches[i] = memdb.NewWatchSet()
	}
	if _, err := state.EvalByID(watches[0], eval1.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalByID(watches[1], eval2.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalsByJob(watches[2], eval1.Namespace, eval1.JobID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.EvalsByJob(watches[3], eval2.Namespace, eval2.JobID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocByID(watches[4], alloc1.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocByID(watches[5], alloc2.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[6], alloc1.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[7], alloc2.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[8], alloc1.Namespace, alloc1.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[9], alloc2.Namespace, alloc2.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[10], alloc1.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[11], alloc2.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	state.UpsertJobSummary(900, mock.JobSummary(eval1.JobID))
	state.UpsertJobSummary(901, mock.JobSummary(eval2.JobID))
	state.UpsertJobSummary(902, mock.JobSummary(alloc1.JobID))
	state.UpsertJobSummary(903, mock.JobSummary(alloc2.JobID))
	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval1, eval2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc1, alloc2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = state.DeleteEval(1002, []string{eval1.ID, eval2.ID}, []string{alloc1.ID, alloc2.ID})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i, ws := range watches {
		if !watchFired(ws) {
			t.Fatalf("bad %d", i)
		}
	}

	ws := memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", eval1, out)
	}

	out, err = state.EvalByID(ws, eval2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", eval1, out)
	}

	outA, err := state.AllocByID(ws, alloc1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", alloc1, outA)
	}

	outA, err = state.AllocByID(ws, alloc2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out != nil {
		t.Fatalf("bad: %#v %#v", alloc1, outA)
	}

	index, err := state.Index("evals")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1002 {
		t.Fatalf("bad: %d", index)
	}

	index, err = state.Index("allocs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1002 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteEval_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	parent := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
		t.Fatalf("err: %v", err)
	}

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
		t.Fatalf("err: %v", err)
	}

	eval1 := mock.Eval()
	eval1.JobID = child.ID
	alloc1 := mock.Alloc()
	alloc1.JobID = child.ID

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval1})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc1})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create watchsets so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.DeleteEval(1002, []string{eval1.ID}, []string{alloc1.ID})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_EvalsByJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	eval1 := mock.Eval()
	eval2 := mock.Eval()
	eval2.JobID = eval1.JobID
	eval3 := mock.Eval()
	evals := []*structs.Evaluation{eval1, eval2}

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, evals)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval3})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.EvalsByJob(ws, eval1.Namespace, eval1.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	sort.Sort(EvalIDSort(evals))
	sort.Sort(EvalIDSort(out))

	if !reflect.DeepEqual(evals, out) {
		t.Fatalf("bad: %#v %#v", evals, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Evals(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var evals []*structs.Evaluation

	for i := 0; i < 10; i++ {
		eval := mock.Eval()
		evals = append(evals, eval)

		err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000+uint64(i), []*structs.Evaluation{eval})
		if err != nil {
			t.Fatalf("err: %v", err)
		}
	}

	ws := memdb.NewWatchSet()
	iter, err := state.Evals(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*structs.Evaluation
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Evaluation))
	}

	sort.Sort(EvalIDSort(evals))
	sort.Sort(EvalIDSort(out))

	if !reflect.DeepEqual(evals, out) {
		t.Fatalf("bad: %#v %#v", evals, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_EvalsByIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var evals []*structs.Evaluation

	ids := []string{
		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
	}
	for i := 0; i < 9; i++ {
		eval := mock.Eval()
		eval.ID = ids[i]
		evals = append(evals, eval)
	}

	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, evals)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	iter, err := state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	gatherEvals := func(iter memdb.ResultIterator) []*structs.Evaluation {
		var evals []*structs.Evaluation
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			evals = append(evals, raw.(*structs.Evaluation))
		}
		return evals
	}

	out := gatherEvals(iter)
	if len(out) != 5 {
		t.Fatalf("bad: expected five evaluations, got: %#v", out)
	}

	sort.Sort(EvalIDSort(evals))

	for index, eval := range out {
		if ids[index] != eval.ID {
			t.Fatalf("bad: got unexpected id: %s", eval.ID)
		}
	}

	iter, err = state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	out = gatherEvals(iter)
	if len(out) != 0 {
		t.Fatalf("bad: unexpected zero evaluations, got: %#v", out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreEval(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	eval := mock.Eval()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.EvalRestore(eval)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.EvalByID(ws, eval.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, eval) {
		t.Fatalf("Bad: %#v %#v", out, eval)
	}
}

func TestStateStore_UpdateAllocsFromClient(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	parent := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
		t.Fatalf("err: %v", err)
	}

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc := mock.Alloc()
	alloc.JobID = child.ID
	alloc.Job = child

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 0 || summary.Children.Running != 1 || summary.Children.Dead != 0 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}

	// Create watchsets so we can test that update fires the watch
	ws = memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create the delta updates
	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStateRunning}}
	update := &structs.Allocation{
		ID:           alloc.ID,
		ClientStatus: structs.AllocClientStatusComplete,
		TaskStates:   ts,
		JobID:        alloc.JobID,
		TaskGroup:    alloc.TaskGroup,
	}
	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	summary, err = state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if summary == nil {
		t.Fatalf("nil summary")
	}
	if summary.JobID != parent.ID {
		t.Fatalf("bad summary id: %v", parent.ID)
	}
	if summary.Children == nil {
		t.Fatalf("nil children summary")
	}
	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
		t.Fatalf("bad children summary: %v", summary.Children)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpdateAllocsFromClient_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc1 := mock.Alloc()
	alloc2 := mock.Alloc()

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc1.Job); err != nil {
		t.Fatalf("err: %v", err)
	}
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc2.Job); err != nil {
		t.Fatalf("err: %v", err)
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc1, alloc2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create watchsets so we can test that update fires the watch
	watches := make([]memdb.WatchSet, 8)
	for i := 0; i < 8; i++ {
		watches[i] = memdb.NewWatchSet()
	}
	if _, err := state.AllocByID(watches[0], alloc1.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocByID(watches[1], alloc2.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[2], alloc1.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[3], alloc2.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[4], alloc1.Namespace, alloc1.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[5], alloc2.Namespace, alloc2.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[6], alloc1.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[7], alloc2.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create the delta updates
	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
	update := &structs.Allocation{
		ID:           alloc1.ID,
		ClientStatus: structs.AllocClientStatusFailed,
		TaskStates:   ts,
		JobID:        alloc1.JobID,
		TaskGroup:    alloc1.TaskGroup,
	}
	update2 := &structs.Allocation{
		ID:           alloc2.ID,
		ClientStatus: structs.AllocClientStatusRunning,
		TaskStates:   ts,
		JobID:        alloc2.JobID,
		TaskGroup:    alloc2.TaskGroup,
	}

	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update, update2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i, ws := range watches {
		if !watchFired(ws) {
			t.Fatalf("bad %d", i)
		}
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc1.CreateIndex = 1000
	alloc1.ModifyIndex = 1001
	alloc1.TaskStates = ts
	alloc1.ClientStatus = structs.AllocClientStatusFailed
	if !reflect.DeepEqual(alloc1, out) {
		t.Fatalf("bad: %#v %#v", alloc1, out)
	}

	out, err = state.AllocByID(ws, alloc2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc2.ModifyIndex = 1000
	alloc2.ModifyIndex = 1001
	alloc2.ClientStatus = structs.AllocClientStatusRunning
	alloc2.TaskStates = ts
	if !reflect.DeepEqual(alloc2, out) {
		t.Fatalf("bad: %#v %#v", alloc2, out)
	}

	index, err := state.Index("allocs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	// Ensure summaries have been updated
	summary, err := state.JobSummaryByID(ws, alloc1.Namespace, alloc1.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	tgSummary := summary.Summary["web"]
	if tgSummary.Failed != 1 {
		t.Fatalf("expected failed: %v, actual: %v, summary: %#v", 1, tgSummary.Failed, tgSummary)
	}

	summary2, err := state.JobSummaryByID(ws, alloc2.Namespace, alloc2.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	tgSummary2 := summary2.Summary["web"]
	if tgSummary2.Running != 1 {
		t.Fatalf("expected running: %v, actual: %v", 1, tgSummary2.Running)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpdateMultipleAllocsFromClient(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
		t.Fatalf("err: %v", err)
	}
	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create the delta updates
	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
	update := &structs.Allocation{
		ID:           alloc.ID,
		ClientStatus: structs.AllocClientStatusRunning,
		TaskStates:   ts,
		JobID:        alloc.JobID,
		TaskGroup:    alloc.TaskGroup,
	}
	update2 := &structs.Allocation{
		ID:           alloc.ID,
		ClientStatus: structs.AllocClientStatusPending,
		TaskStates:   ts,
		JobID:        alloc.JobID,
		TaskGroup:    alloc.TaskGroup,
	}

	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update, update2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc.CreateIndex = 1000
	alloc.ModifyIndex = 1001
	alloc.TaskStates = ts
	alloc.ClientStatus = structs.AllocClientStatusPending
	if !reflect.DeepEqual(alloc, out) {
		t.Fatalf("bad: %#v , actual:%#v", alloc, out)
	}

	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
	expectedSummary := &structs.JobSummary{
		JobID:     alloc.JobID,
		Namespace: alloc.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Starting: 1,
			},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: 999,
		ModifyIndex: 1001,
	}
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !reflect.DeepEqual(summary, expectedSummary) {
		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
	}
}

func TestStateStore_UpdateAllocsFromClient_Deployment(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	alloc := mock.Alloc()
	now := time.Now()
	alloc.CreateTime = now.UnixNano()
	pdeadline := 5 * time.Minute
	deployment := mock.Deployment()
	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
	alloc.DeploymentID = deployment.ID

	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
	require.Nil(state.UpsertDeployment(1000, deployment))
	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}))

	healthy := now.Add(time.Second)
	update := &structs.Allocation{
		ID:           alloc.ID,
		ClientStatus: structs.AllocClientStatusRunning,
		JobID:        alloc.JobID,
		TaskGroup:    alloc.TaskGroup,
		DeploymentStatus: &structs.AllocDeploymentStatus{
			Healthy:   helper.BoolToPtr(true),
			Timestamp: healthy,
		},
	}
	require.Nil(state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update}))

	// Check that the deployment state was updated because the healthy
	// deployment
	dout, err := state.DeploymentByID(nil, deployment.ID)
	require.Nil(err)
	require.NotNil(dout)
	require.Len(dout.TaskGroups, 1)
	dstate := dout.TaskGroups[alloc.TaskGroup]
	require.NotNil(dstate)
	require.Equal(1, dstate.PlacedAllocs)
	require.True(healthy.Add(pdeadline).Equal(dstate.RequireProgressBy))
}

// This tests that the deployment state is merged correctly
func TestStateStore_UpdateAllocsFromClient_DeploymentStateMerges(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	alloc := mock.Alloc()
	now := time.Now()
	alloc.CreateTime = now.UnixNano()
	pdeadline := 5 * time.Minute
	deployment := mock.Deployment()
	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
	alloc.DeploymentID = deployment.ID
	alloc.DeploymentStatus = &structs.AllocDeploymentStatus{
		Canary: true,
	}

	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
	require.Nil(state.UpsertDeployment(1000, deployment))
	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}))

	update := &structs.Allocation{
		ID:           alloc.ID,
		ClientStatus: structs.AllocClientStatusRunning,
		JobID:        alloc.JobID,
		TaskGroup:    alloc.TaskGroup,
		DeploymentStatus: &structs.AllocDeploymentStatus{
			Healthy: helper.BoolToPtr(true),
			Canary:  false,
		},
	}
	require.Nil(state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update}))

	// Check that the merging of the deployment status was correct
	out, err := state.AllocByID(nil, alloc.ID)
	require.Nil(err)
	require.NotNil(out)
	require.True(out.DeploymentStatus.Canary)
	require.NotNil(out.DeploymentStatus.Healthy)
	require.True(*out.DeploymentStatus.Healthy)
}

func TestStateStore_UpsertAlloc_Alloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create watchsets so we can test that update fires the watch
	watches := make([]memdb.WatchSet, 4)
	for i := 0; i < 4; i++ {
		watches[i] = memdb.NewWatchSet()
	}
	if _, err := state.AllocByID(watches[0], alloc.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[1], alloc.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[2], alloc.Namespace, alloc.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[3], alloc.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i, ws := range watches {
		if !watchFired(ws) {
			t.Fatalf("bad %d", i)
		}
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(alloc, out) {
		t.Fatalf("bad: %#v %#v", alloc, out)
	}

	index, err := state.Index("allocs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	tgSummary, ok := summary.Summary["web"]
	if !ok {
		t.Fatalf("no summary for task group web")
	}
	if tgSummary.Starting != 1 {
		t.Fatalf("expected queued: %v, actual: %v", 1, tgSummary.Starting)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertAlloc_Deployment(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	alloc := mock.Alloc()
	now := time.Now()
	alloc.CreateTime = now.UnixNano()
	alloc.ModifyTime = now.UnixNano()
	pdeadline := 5 * time.Minute
	deployment := mock.Deployment()
	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
	alloc.DeploymentID = deployment.ID

	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
	require.Nil(state.UpsertDeployment(1000, deployment))

	// Create a watch set so we can test that update fires the watch
	ws := memdb.NewWatchSet()
	require.Nil(state.AllocsByDeployment(ws, alloc.DeploymentID))

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc})
	require.Nil(err)

	if !watchFired(ws) {
		t.Fatalf("watch not fired")
	}

	ws = memdb.NewWatchSet()
	allocs, err := state.AllocsByDeployment(ws, alloc.DeploymentID)
	require.Nil(err)
	require.Len(allocs, 1)
	require.EqualValues(alloc, allocs[0])

	index, err := state.Index("allocs")
	require.Nil(err)
	require.EqualValues(1001, index)
	if watchFired(ws) {
		t.Fatalf("bad")
	}

	// Check that the deployment state was updated
	dout, err := state.DeploymentByID(nil, deployment.ID)
	require.Nil(err)
	require.NotNil(dout)
	require.Len(dout.TaskGroups, 1)
	dstate := dout.TaskGroups[alloc.TaskGroup]
	require.NotNil(dstate)
	require.Equal(1, dstate.PlacedAllocs)
	require.True(now.Add(pdeadline).Equal(dstate.RequireProgressBy))
}

// Testing to ensure we keep issue
// https://github.com/hashicorp/nomad/issues/2583 fixed
func TestStateStore_UpsertAlloc_No_Job(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	alloc.Job = nil

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 999, []*structs.Allocation{alloc})
	if err == nil || !strings.Contains(err.Error(), "without a job") {
		t.Fatalf("expect err: %v", err)
	}
}

func TestStateStore_UpsertAlloc_ChildJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	parent := mock.Job()
	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 998, parent))

	child := mock.Job()
	child.Status = ""
	child.ParentID = parent.ID

	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 999, child))

	alloc := mock.Alloc()
	alloc.JobID = child.ID
	alloc.Job = child

	// Create watchsets so we can test that delete fires the watch
	ws := memdb.NewWatchSet()
	_, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	require.NoError(t, err)

	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	require.NoError(t, err)

	require.True(t, watchFired(ws))

	ws = memdb.NewWatchSet()
	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
	require.NoError(t, err)
	require.NotNil(t, summary)

	require.Equal(t, parent.ID, summary.JobID)
	require.NotNil(t, summary.Children)

	require.Equal(t, int64(0), summary.Children.Pending)
	require.Equal(t, int64(1), summary.Children.Running)
	require.Equal(t, int64(0), summary.Children.Dead)

	require.False(t, watchFired(ws))
}

func TestStateStore_UpdateAlloc_Alloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
		t.Fatalf("err: %v", err)
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	tgSummary := summary.Summary["web"]
	if tgSummary.Starting != 1 {
		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
	}

	alloc2 := mock.Alloc()
	alloc2.ID = alloc.ID
	alloc2.NodeID = alloc.NodeID + ".new"
	state.UpsertJobSummary(1001, mock.JobSummary(alloc2.JobID))

	// Create watchsets so we can test that update fires the watch
	watches := make([]memdb.WatchSet, 4)
	for i := 0; i < 4; i++ {
		watches[i] = memdb.NewWatchSet()
	}
	if _, err := state.AllocByID(watches[0], alloc2.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByEval(watches[1], alloc2.EvalID); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByJob(watches[2], alloc2.Namespace, alloc2.JobID, false); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if _, err := state.AllocsByNode(watches[3], alloc2.NodeID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	for i, ws := range watches {
		if !watchFired(ws) {
			t.Fatalf("bad %d", i)
		}
	}

	ws = memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(alloc2, out) {
		t.Fatalf("bad: %#v %#v", alloc2, out)
	}

	if out.CreateIndex != 1000 {
		t.Fatalf("bad: %#v", out)
	}
	if out.ModifyIndex != 1002 {
		t.Fatalf("bad: %#v", out)
	}

	index, err := state.Index("allocs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1002 {
		t.Fatalf("bad: %d", index)
	}

	// Ensure that summary hasb't changed
	summary, err = state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	tgSummary = summary.Summary["web"]
	if tgSummary.Starting != 1 {
		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

// This test ensures that the state store will mark the clients status as lost
// when set rather than preferring the existing status.
func TestStateStore_UpdateAlloc_Lost(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	alloc.ClientStatus = "foo"

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
		t.Fatalf("err: %v", err)
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc2 := new(structs.Allocation)
	*alloc2 = *alloc
	alloc2.ClientStatus = structs.AllocClientStatusLost
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc2}); err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out.ClientStatus != structs.AllocClientStatusLost {
		t.Fatalf("bad: %#v", out)
	}
}

// This test ensures an allocation can be updated when there is no job
// associated with it. This will happen when a job is stopped by an user which
// has non-terminal allocations on clients
func TestStateStore_UpdateAlloc_NoJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	// Upsert a job
	state.UpsertJobSummary(998, mock.JobSummary(alloc.JobID))
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
		t.Fatalf("err: %v", err)
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.DeleteJob(1001, alloc.Namespace, alloc.JobID); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Update the desired state of the allocation to stop
	allocCopy := alloc.Copy()
	allocCopy.DesiredStatus = structs.AllocDesiredStatusStop
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{allocCopy}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Update the client state of the allocation to complete
	allocCopy1 := allocCopy.Copy()
	allocCopy1.ClientStatus = structs.AllocClientStatusComplete
	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1003, []*structs.Allocation{allocCopy1}); err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, _ := state.AllocByID(ws, alloc.ID)
	// Update the modify index of the alloc before comparing
	allocCopy1.ModifyIndex = 1003
	if !reflect.DeepEqual(out, allocCopy1) {
		t.Fatalf("expected: %#v \n actual: %#v", allocCopy1, out)
	}
}

func TestStateStore_UpdateAllocDesiredTransition(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	alloc := mock.Alloc()

	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}))

	t1 := &structs.DesiredTransition{
		Migrate: helper.BoolToPtr(true),
	}
	t2 := &structs.DesiredTransition{
		Migrate: helper.BoolToPtr(false),
	}
	eval := &structs.Evaluation{
		ID:             uuid.Generate(),
		Namespace:      alloc.Namespace,
		Priority:       alloc.Job.Priority,
		Type:           alloc.Job.Type,
		TriggeredBy:    structs.EvalTriggerNodeDrain,
		JobID:          alloc.Job.ID,
		JobModifyIndex: alloc.Job.ModifyIndex,
		Status:         structs.EvalStatusPending,
	}
	evals := []*structs.Evaluation{eval}

	m := map[string]*structs.DesiredTransition{alloc.ID: t1}
	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1001, m, evals))

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	require.Nil(err)
	require.NotNil(out.DesiredTransition.Migrate)
	require.True(*out.DesiredTransition.Migrate)
	require.EqualValues(1000, out.CreateIndex)
	require.EqualValues(1001, out.ModifyIndex)

	index, err := state.Index("allocs")
	require.Nil(err)
	require.EqualValues(1001, index)

	// Check the eval is created
	eout, err := state.EvalByID(nil, eval.ID)
	require.Nil(err)
	require.NotNil(eout)

	m = map[string]*structs.DesiredTransition{alloc.ID: t2}
	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1002, m, evals))

	ws = memdb.NewWatchSet()
	out, err = state.AllocByID(ws, alloc.ID)
	require.Nil(err)
	require.NotNil(out.DesiredTransition.Migrate)
	require.False(*out.DesiredTransition.Migrate)
	require.EqualValues(1000, out.CreateIndex)
	require.EqualValues(1002, out.ModifyIndex)

	index, err = state.Index("allocs")
	require.Nil(err)
	require.EqualValues(1002, index)

	// Try with a bogus alloc id
	m = map[string]*structs.DesiredTransition{uuid.Generate(): t2}
	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1003, m, evals))
}

func TestStateStore_JobSummary(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Add a job
	job := mock.Job()
	state.UpsertJob(structs.MsgTypeTestSetup, 900, job)

	// Get the job back
	ws := memdb.NewWatchSet()
	outJob, _ := state.JobByID(ws, job.Namespace, job.ID)
	if outJob.CreateIndex != 900 {
		t.Fatalf("bad create index: %v", outJob.CreateIndex)
	}
	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if summary.CreateIndex != 900 {
		t.Fatalf("bad create index: %v", summary.CreateIndex)
	}

	// Upsert an allocation
	alloc := mock.Alloc()
	alloc.JobID = job.ID
	alloc.Job = job
	state.UpsertAllocs(structs.MsgTypeTestSetup, 910, []*structs.Allocation{alloc})

	// Update the alloc from client
	alloc1 := alloc.Copy()
	alloc1.ClientStatus = structs.AllocClientStatusPending
	alloc1.DesiredStatus = ""
	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 920, []*structs.Allocation{alloc})

	alloc3 := alloc.Copy()
	alloc3.ClientStatus = structs.AllocClientStatusRunning
	alloc3.DesiredStatus = ""
	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 930, []*structs.Allocation{alloc3})

	// Upsert the alloc
	alloc4 := alloc.Copy()
	alloc4.ClientStatus = structs.AllocClientStatusPending
	alloc4.DesiredStatus = structs.AllocDesiredStatusRun
	state.UpsertAllocs(structs.MsgTypeTestSetup, 950, []*structs.Allocation{alloc4})

	// Again upsert the alloc
	alloc5 := alloc.Copy()
	alloc5.ClientStatus = structs.AllocClientStatusPending
	alloc5.DesiredStatus = structs.AllocDesiredStatusRun
	state.UpsertAllocs(structs.MsgTypeTestSetup, 970, []*structs.Allocation{alloc5})

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	expectedSummary := structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Running: 1,
			},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: 900,
		ModifyIndex: 930,
	}

	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	if !reflect.DeepEqual(&expectedSummary, summary) {
		t.Fatalf("expected: %#v, actual: %v", expectedSummary, summary)
	}

	// De-register the job.
	state.DeleteJob(980, job.Namespace, job.ID)

	// Shouldn't have any effect on the summary
	alloc6 := alloc.Copy()
	alloc6.ClientStatus = structs.AllocClientStatusRunning
	alloc6.DesiredStatus = ""
	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 990, []*structs.Allocation{alloc6})

	// We shouldn't have any summary at this point
	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	if summary != nil {
		t.Fatalf("expected nil, actual: %#v", summary)
	}

	// Re-register the same job
	job1 := mock.Job()
	job1.ID = job.ID
	state.UpsertJob(structs.MsgTypeTestSetup, 1000, job1)
	outJob2, _ := state.JobByID(ws, job1.Namespace, job1.ID)
	if outJob2.CreateIndex != 1000 {
		t.Fatalf("bad create index: %v", outJob2.CreateIndex)
	}
	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
	if summary.CreateIndex != 1000 {
		t.Fatalf("bad create index: %v", summary.CreateIndex)
	}

	// Upsert an allocation
	alloc7 := alloc.Copy()
	alloc7.JobID = outJob.ID
	alloc7.Job = outJob
	alloc7.ClientStatus = structs.AllocClientStatusComplete
	alloc7.DesiredStatus = structs.AllocDesiredStatusRun
	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1020, []*structs.Allocation{alloc7})

	expectedSummary = structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: 1000,
		ModifyIndex: 1000,
	}

	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
	if !reflect.DeepEqual(&expectedSummary, summary) {
		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
	}
}

func TestStateStore_ReconcileJobSummary(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create an alloc
	alloc := mock.Alloc()

	// Add another task group to the job
	tg2 := alloc.Job.TaskGroups[0].Copy()
	tg2.Name = "db"
	alloc.Job.TaskGroups = append(alloc.Job.TaskGroups, tg2)
	state.UpsertJob(structs.MsgTypeTestSetup, 100, alloc.Job)

	// Create one more alloc for the db task group
	alloc2 := mock.Alloc()
	alloc2.TaskGroup = "db"
	alloc2.JobID = alloc.JobID
	alloc2.Job = alloc.Job

	// Upserts the alloc
	state.UpsertAllocs(structs.MsgTypeTestSetup, 110, []*structs.Allocation{alloc, alloc2})

	// Change the state of the first alloc to running
	alloc3 := alloc.Copy()
	alloc3.ClientStatus = structs.AllocClientStatusRunning
	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 120, []*structs.Allocation{alloc3})

	//Add some more allocs to the second tg
	alloc4 := mock.Alloc()
	alloc4.JobID = alloc.JobID
	alloc4.Job = alloc.Job
	alloc4.TaskGroup = "db"
	alloc5 := alloc4.Copy()
	alloc5.ClientStatus = structs.AllocClientStatusRunning

	alloc6 := mock.Alloc()
	alloc6.JobID = alloc.JobID
	alloc6.Job = alloc.Job
	alloc6.TaskGroup = "db"
	alloc7 := alloc6.Copy()
	alloc7.ClientStatus = structs.AllocClientStatusComplete

	alloc8 := mock.Alloc()
	alloc8.JobID = alloc.JobID
	alloc8.Job = alloc.Job
	alloc8.TaskGroup = "db"
	alloc9 := alloc8.Copy()
	alloc9.ClientStatus = structs.AllocClientStatusFailed

	alloc10 := mock.Alloc()
	alloc10.JobID = alloc.JobID
	alloc10.Job = alloc.Job
	alloc10.TaskGroup = "db"
	alloc11 := alloc10.Copy()
	alloc11.ClientStatus = structs.AllocClientStatusLost

	state.UpsertAllocs(structs.MsgTypeTestSetup, 130, []*structs.Allocation{alloc4, alloc6, alloc8, alloc10})

	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 150, []*structs.Allocation{alloc5, alloc7, alloc9, alloc11})

	// DeleteJobSummary is a helper method and doesn't modify the indexes table
	state.DeleteJobSummary(130, alloc.Namespace, alloc.Job.ID)

	state.ReconcileJobSummaries(120)

	ws := memdb.NewWatchSet()
	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
	expectedSummary := structs.JobSummary{
		JobID:     alloc.Job.ID,
		Namespace: alloc.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Running: 1,
			},
			"db": {
				Starting: 1,
				Running:  1,
				Failed:   1,
				Complete: 1,
				Lost:     1,
			},
		},
		CreateIndex: 100,
		ModifyIndex: 120,
	}
	if !reflect.DeepEqual(&expectedSummary, summary) {
		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
	}
}

func TestStateStore_ReconcileParentJobSummary(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	// Add a node
	node := mock.Node()
	state.UpsertNode(structs.MsgTypeTestSetup, 80, node)

	// Make a parameterized job
	job1 := mock.BatchJob()
	job1.ID = "test"
	job1.ParameterizedJob = &structs.ParameterizedJobConfig{
		Payload: "random",
	}
	job1.TaskGroups[0].Count = 1
	state.UpsertJob(structs.MsgTypeTestSetup, 100, job1)

	// Make a child job
	childJob := job1.Copy()
	childJob.ID = job1.ID + "dispatch-23423423"
	childJob.ParentID = job1.ID
	childJob.Dispatched = true
	childJob.Status = structs.JobStatusRunning

	// Make some allocs for child job
	alloc := mock.Alloc()
	alloc.NodeID = node.ID
	alloc.Job = childJob
	alloc.JobID = childJob.ID
	alloc.ClientStatus = structs.AllocClientStatusRunning

	alloc2 := mock.Alloc()
	alloc2.NodeID = node.ID
	alloc2.Job = childJob
	alloc2.JobID = childJob.ID
	alloc2.ClientStatus = structs.AllocClientStatusFailed

	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 110, childJob))
	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 111, []*structs.Allocation{alloc, alloc2}))

	// Make the summary incorrect in the state store
	summary, err := state.JobSummaryByID(nil, job1.Namespace, job1.ID)
	require.Nil(err)

	summary.Children = nil
	summary.Summary = make(map[string]structs.TaskGroupSummary)
	summary.Summary["web"] = structs.TaskGroupSummary{
		Queued: 1,
	}

	// Delete the child job summary
	state.DeleteJobSummary(125, childJob.Namespace, childJob.ID)

	state.ReconcileJobSummaries(120)

	ws := memdb.NewWatchSet()

	// Verify parent summary is corrected
	summary, _ = state.JobSummaryByID(ws, alloc.Namespace, job1.ID)
	expectedSummary := structs.JobSummary{
		JobID:     job1.ID,
		Namespace: job1.Namespace,
		Summary:   make(map[string]structs.TaskGroupSummary),
		Children: &structs.JobChildrenSummary{
			Running: 1,
		},
		CreateIndex: 100,
		ModifyIndex: 120,
	}
	require.Equal(&expectedSummary, summary)

	// Verify child job summary is also correct
	childSummary, _ := state.JobSummaryByID(ws, childJob.Namespace, childJob.ID)
	expectedChildSummary := structs.JobSummary{
		JobID:     childJob.ID,
		Namespace: childJob.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Running: 1,
				Failed:  1,
			},
		},
		CreateIndex: 110,
		ModifyIndex: 120,
	}
	require.Equal(&expectedChildSummary, childSummary)
}

func TestStateStore_UpdateAlloc_JobNotPresent(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	alloc := mock.Alloc()
	state.UpsertJob(structs.MsgTypeTestSetup, 100, alloc.Job)
	state.UpsertAllocs(structs.MsgTypeTestSetup, 200, []*structs.Allocation{alloc})

	// Delete the job
	state.DeleteJob(300, alloc.Namespace, alloc.Job.ID)

	// Update the alloc
	alloc1 := alloc.Copy()
	alloc1.ClientStatus = structs.AllocClientStatusRunning

	// Updating allocation should not throw any error
	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 400, []*structs.Allocation{alloc1}); err != nil {
		t.Fatalf("expect err: %v", err)
	}

	// Re-Register the job
	state.UpsertJob(structs.MsgTypeTestSetup, 500, alloc.Job)

	// Update the alloc again
	alloc2 := alloc.Copy()
	alloc2.ClientStatus = structs.AllocClientStatusComplete
	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 400, []*structs.Allocation{alloc1}); err != nil {
		t.Fatalf("expect err: %v", err)
	}

	// Job Summary of the newly registered job shouldn't account for the
	// allocation update for the older job
	expectedSummary := structs.JobSummary{
		JobID:     alloc1.JobID,
		Namespace: alloc1.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: 500,
		ModifyIndex: 500,
	}

	ws := memdb.NewWatchSet()
	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
	if !reflect.DeepEqual(&expectedSummary, summary) {
		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
	}
}

func TestStateStore_EvictAlloc_Alloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	alloc2 := new(structs.Allocation)
	*alloc2 = *alloc
	alloc2.DesiredStatus = structs.AllocDesiredStatusEvict
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if out.DesiredStatus != structs.AllocDesiredStatusEvict {
		t.Fatalf("bad: %#v %#v", alloc, out)
	}

	index, err := state.Index("allocs")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}
}

func TestStateStore_AllocsByNode(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation

	for i := 0; i < 10; i++ {
		alloc := mock.Alloc()
		alloc.NodeID = "foo"
		allocs = append(allocs, alloc)
	}

	for idx, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocsByNode(ws, "foo")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	sort.Sort(AllocIDSort(allocs))
	sort.Sort(AllocIDSort(out))

	if !reflect.DeepEqual(allocs, out) {
		t.Fatalf("bad: %#v %#v", allocs, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_AllocsByNodeTerminal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs, term, nonterm []*structs.Allocation

	for i := 0; i < 10; i++ {
		alloc := mock.Alloc()
		alloc.NodeID = "foo"
		if i%2 == 0 {
			alloc.DesiredStatus = structs.AllocDesiredStatusStop
			term = append(term, alloc)
		} else {
			nonterm = append(nonterm, alloc)
		}
		allocs = append(allocs, alloc)
	}

	for idx, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Verify the terminal allocs
	ws := memdb.NewWatchSet()
	out, err := state.AllocsByNodeTerminal(ws, "foo", true)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	sort.Sort(AllocIDSort(term))
	sort.Sort(AllocIDSort(out))

	if !reflect.DeepEqual(term, out) {
		t.Fatalf("bad: %#v %#v", term, out)
	}

	// Verify the non-terminal allocs
	out, err = state.AllocsByNodeTerminal(ws, "foo", false)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	sort.Sort(AllocIDSort(nonterm))
	sort.Sort(AllocIDSort(out))

	if !reflect.DeepEqual(nonterm, out) {
		t.Fatalf("bad: %#v %#v", nonterm, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_AllocsByJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation

	for i := 0; i < 10; i++ {
		alloc := mock.Alloc()
		alloc.JobID = "foo"
		allocs = append(allocs, alloc)
	}

	for i, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocsByJob(ws, mock.Alloc().Namespace, "foo", false)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	sort.Sort(AllocIDSort(allocs))
	sort.Sort(AllocIDSort(out))

	if !reflect.DeepEqual(allocs, out) {
		t.Fatalf("bad: %#v %#v", allocs, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_AllocsForRegisteredJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation
	var allocs1 []*structs.Allocation

	job := mock.Job()
	job.ID = "foo"
	state.UpsertJob(structs.MsgTypeTestSetup, 100, job)
	for i := 0; i < 3; i++ {
		alloc := mock.Alloc()
		alloc.Job = job
		alloc.JobID = job.ID
		allocs = append(allocs, alloc)
	}
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 200, allocs); err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.DeleteJob(250, job.Namespace, job.ID); err != nil {
		t.Fatalf("err: %v", err)
	}

	job1 := mock.Job()
	job1.ID = "foo"
	job1.CreateIndex = 50
	state.UpsertJob(structs.MsgTypeTestSetup, 300, job1)
	for i := 0; i < 4; i++ {
		alloc := mock.Alloc()
		alloc.Job = job1
		alloc.JobID = job1.ID
		allocs1 = append(allocs1, alloc)
	}

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs1); err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, true)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	expected := len(allocs) + len(allocs1)
	if len(out) != expected {
		t.Fatalf("expected: %v, actual: %v", expected, len(out))
	}

	out1, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, false)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	expected = len(allocs1)
	if len(out1) != expected {
		t.Fatalf("expected: %v, actual: %v", expected, len(out1))
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_AllocsByIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation

	ids := []string{
		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
	}
	for i := 0; i < 9; i++ {
		alloc := mock.Alloc()
		alloc.ID = ids[i]
		allocs = append(allocs, alloc)
	}

	for i, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	iter, err := state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	gatherAllocs := func(iter memdb.ResultIterator) []*structs.Allocation {
		var allocs []*structs.Allocation
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			allocs = append(allocs, raw.(*structs.Allocation))
		}
		return allocs
	}

	out := gatherAllocs(iter)
	if len(out) != 5 {
		t.Fatalf("bad: expected five allocations, got: %#v", out)
	}

	sort.Sort(AllocIDSort(allocs))

	for index, alloc := range out {
		if ids[index] != alloc.ID {
			t.Fatalf("bad: got unexpected id: %s", alloc.ID)
		}
	}

	iter, err = state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	out = gatherAllocs(iter)
	if len(out) != 0 {
		t.Fatalf("bad: unexpected zero allocations, got: %#v", out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Allocs(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation

	for i := 0; i < 10; i++ {
		alloc := mock.Alloc()
		allocs = append(allocs, alloc)
	}
	for i, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
	}

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	iter, err := state.Allocs(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	var out []*structs.Allocation
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Allocation))
	}

	sort.Sort(AllocIDSort(allocs))
	sort.Sort(AllocIDSort(out))

	if !reflect.DeepEqual(allocs, out) {
		t.Fatalf("bad: %#v %#v", allocs, out)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_Allocs_PrevAlloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	var allocs []*structs.Allocation

	require := require.New(t)
	for i := 0; i < 5; i++ {
		alloc := mock.Alloc()
		allocs = append(allocs, alloc)
	}
	for i, alloc := range allocs {
		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
	}
	// Set some previous alloc ids
	allocs[1].PreviousAllocation = allocs[0].ID
	allocs[2].PreviousAllocation = allocs[1].ID

	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
	require.Nil(err)

	ws := memdb.NewWatchSet()
	iter, err := state.Allocs(ws)
	require.Nil(err)

	var out []*structs.Allocation
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		out = append(out, raw.(*structs.Allocation))
	}

	// Set expected NextAllocation fields
	allocs[0].NextAllocation = allocs[1].ID
	allocs[1].NextAllocation = allocs[2].ID

	sort.Sort(AllocIDSort(allocs))
	sort.Sort(AllocIDSort(out))

	require.Equal(allocs, out)
	require.False(watchFired(ws))

	// Insert another alloc, verify index of previous alloc also got updated
	alloc := mock.Alloc()
	alloc.PreviousAllocation = allocs[0].ID
	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc})
	require.Nil(err)
	alloc0, err := state.AllocByID(nil, allocs[0].ID)
	require.Nil(err)
	require.Equal(alloc0.ModifyIndex, uint64(1001))
}

func TestStateStore_RestoreAlloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.AllocRestore(alloc)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.AllocByID(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, alloc) {
		t.Fatalf("Bad: %#v %#v", out, alloc)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_SetJobStatus_ForceStatus(t *testing.T) {
	t.Parallel()

	index := uint64(0)
	state := testStateStore(t)
	txn := state.db.WriteTxn(index)

	// Create and insert a mock job.
	job := mock.Job()
	job.Status = ""
	job.ModifyIndex = index
	if err := txn.Insert("jobs", job); err != nil {
		t.Fatalf("job insert failed: %v", err)
	}

	exp := "foobar"
	index = uint64(1000)
	if err := state.setJobStatus(index, txn, job, false, exp); err != nil {
		t.Fatalf("setJobStatus() failed: %v", err)
	}

	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("job lookup failed: %v", err)
	}
	updated := i.(*structs.Job)

	if updated.Status != exp {
		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, exp)
	}

	if updated.ModifyIndex != index {
		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
	}
}

func TestStateStore_SetJobStatus_NoOp(t *testing.T) {
	t.Parallel()

	index := uint64(0)
	state := testStateStore(t)
	txn := state.db.WriteTxn(index)

	// Create and insert a mock job that should be pending.
	job := mock.Job()
	job.Status = structs.JobStatusPending
	job.ModifyIndex = 10
	if err := txn.Insert("jobs", job); err != nil {
		t.Fatalf("job insert failed: %v", err)
	}

	index = uint64(1000)
	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
		t.Fatalf("setJobStatus() failed: %v", err)
	}

	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("job lookup failed: %v", err)
	}
	updated := i.(*structs.Job)

	if updated.ModifyIndex == index {
		t.Fatalf("setJobStatus() should have been a no-op")
	}
}

func TestStateStore_SetJobStatus(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	txn := state.db.WriteTxn(uint64(0))

	// Create and insert a mock job that should be pending but has an incorrect
	// status.
	job := mock.Job()
	job.Status = "foobar"
	job.ModifyIndex = 10
	if err := txn.Insert("jobs", job); err != nil {
		t.Fatalf("job insert failed: %v", err)
	}

	index := uint64(1000)
	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
		t.Fatalf("setJobStatus() failed: %v", err)
	}

	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("job lookup failed: %v", err)
	}
	updated := i.(*structs.Job)

	if updated.Status != structs.JobStatusPending {
		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, structs.JobStatusPending)
	}

	if updated.ModifyIndex != index {
		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
	}
}

func TestStateStore_GetJobStatus_NoEvalsOrAllocs(t *testing.T) {
	t.Parallel()

	job := mock.Job()
	state := testStateStore(t)
	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusPending {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
	}
}

func TestStateStore_GetJobStatus_NoEvalsOrAllocs_Periodic(t *testing.T) {
	t.Parallel()

	job := mock.PeriodicJob()
	state := testStateStore(t)
	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusRunning {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
	}
}

func TestStateStore_GetJobStatus_NoEvalsOrAllocs_EvalDelete(t *testing.T) {
	t.Parallel()

	job := mock.Job()
	state := testStateStore(t)
	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, true)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusDead {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
	}
}

func TestStateStore_GetJobStatus_DeadEvalsAndAllocs(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	// Create a mock alloc that is dead.
	alloc := mock.Alloc()
	alloc.JobID = job.ID
	alloc.DesiredStatus = structs.AllocDesiredStatusStop
	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a mock eval that is complete
	eval := mock.Eval()
	eval.JobID = job.ID
	eval.Status = structs.EvalStatusComplete
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusDead {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
	}
}

func TestStateStore_GetJobStatus_RunningAlloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	// Create a mock alloc that is running.
	alloc := mock.Alloc()
	alloc.JobID = job.ID
	alloc.DesiredStatus = structs.AllocDesiredStatusRun
	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}); err != nil {
		t.Fatalf("err: %v", err)
	}

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, true)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusRunning {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
	}
}

func TestStateStore_GetJobStatus_PeriodicJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.PeriodicJob()

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusRunning {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
	}

	// Mark it as stopped
	job.Stop = true
	status, err = state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusDead {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
	}
}

func TestStateStore_GetJobStatus_ParameterizedJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()
	job.ParameterizedJob = &structs.ParameterizedJobConfig{}

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusRunning {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
	}

	// Mark it as stopped
	job.Stop = true
	status, err = state.getJobStatus(txn, job, false)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusDead {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
	}
}

func TestStateStore_SetJobStatus_PendingEval(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.Job()

	// Create a mock eval that is pending.
	eval := mock.Eval()
	eval.JobID = job.ID
	eval.Status = structs.EvalStatusPending
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, true)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if status != structs.JobStatusPending {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
	}
}

// TestStateStore_SetJobStatus_SystemJob asserts that system jobs are still
// considered running until explicitly stopped.
func TestStateStore_SetJobStatus_SystemJob(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	job := mock.SystemJob()

	// Create a mock eval that is pending.
	eval := mock.Eval()
	eval.JobID = job.ID
	eval.Type = job.Type
	eval.Status = structs.EvalStatusComplete
	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
		t.Fatalf("err: %v", err)
	}

	txn := state.db.ReadTxn()
	status, err := state.getJobStatus(txn, job, true)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if expected := structs.JobStatusRunning; status != expected {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
	}

	// Stop the job
	job.Stop = true
	status, err = state.getJobStatus(txn, job, true)
	if err != nil {
		t.Fatalf("getJobStatus() failed: %v", err)
	}

	if expected := structs.JobStatusDead; status != expected {
		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
	}
}

func TestStateJobSummary_UpdateJobCount(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	job := alloc.Job
	job.TaskGroups[0].Count = 3

	// Create watchsets so we can test that upsert fires the watch
	ws := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws, job.Namespace, job.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
	expectedSummary := structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Starting: 1,
			},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: 1000,
		ModifyIndex: 1001,
	}
	if !reflect.DeepEqual(summary, &expectedSummary) {
		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
	}

	// Create watchsets so we can test that upsert fires the watch
	ws2 := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws2, job.Namespace, job.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	alloc2 := mock.Alloc()
	alloc2.Job = job
	alloc2.JobID = job.ID

	alloc3 := mock.Alloc()
	alloc3.Job = job
	alloc3.JobID = job.ID

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc2, alloc3}); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws2) {
		t.Fatalf("bad")
	}

	outA, _ := state.AllocByID(ws, alloc3.ID)

	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	expectedSummary = structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Starting: 3,
			},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: job.CreateIndex,
		ModifyIndex: outA.ModifyIndex,
	}
	if !reflect.DeepEqual(summary, &expectedSummary) {
		t.Fatalf("expected summary: %v, actual: %v", expectedSummary, summary)
	}

	// Create watchsets so we can test that upsert fires the watch
	ws3 := memdb.NewWatchSet()
	if _, err := state.JobSummaryByID(ws3, job.Namespace, job.ID); err != nil {
		t.Fatalf("bad: %v", err)
	}

	alloc4 := mock.Alloc()
	alloc4.ID = alloc2.ID
	alloc4.Job = alloc2.Job
	alloc4.JobID = alloc2.JobID
	alloc4.ClientStatus = structs.AllocClientStatusComplete

	alloc5 := mock.Alloc()
	alloc5.ID = alloc3.ID
	alloc5.Job = alloc3.Job
	alloc5.JobID = alloc3.JobID
	alloc5.ClientStatus = structs.AllocClientStatusComplete

	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1004, []*structs.Allocation{alloc4, alloc5}); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws2) {
		t.Fatalf("bad")
	}

	outA, _ = state.AllocByID(ws, alloc5.ID)
	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	expectedSummary = structs.JobSummary{
		JobID:     job.ID,
		Namespace: job.Namespace,
		Summary: map[string]structs.TaskGroupSummary{
			"web": {
				Complete: 2,
				Starting: 1,
			},
		},
		Children:    new(structs.JobChildrenSummary),
		CreateIndex: job.CreateIndex,
		ModifyIndex: outA.ModifyIndex,
	}
	if !reflect.DeepEqual(summary, &expectedSummary) {
		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
	}
}

func TestJobSummary_UpdateClientStatus(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	job := alloc.Job
	job.TaskGroups[0].Count = 3

	alloc2 := mock.Alloc()
	alloc2.Job = job
	alloc2.JobID = job.ID

	alloc3 := mock.Alloc()
	alloc3.Job = job
	alloc3.JobID = job.ID

	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc, alloc2, alloc3}); err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
	if summary.Summary["web"].Starting != 3 {
		t.Fatalf("bad job summary: %v", summary)
	}

	alloc4 := mock.Alloc()
	alloc4.ID = alloc2.ID
	alloc4.Job = alloc2.Job
	alloc4.JobID = alloc2.JobID
	alloc4.ClientStatus = structs.AllocClientStatusComplete

	alloc5 := mock.Alloc()
	alloc5.ID = alloc3.ID
	alloc5.Job = alloc3.Job
	alloc5.JobID = alloc3.JobID
	alloc5.ClientStatus = structs.AllocClientStatusFailed

	alloc6 := mock.Alloc()
	alloc6.ID = alloc.ID
	alloc6.Job = alloc.Job
	alloc6.JobID = alloc.JobID
	alloc6.ClientStatus = structs.AllocClientStatusRunning

	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc4, alloc5, alloc6}); err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	if summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
		t.Fatalf("bad job summary: %v", summary)
	}

	alloc7 := mock.Alloc()
	alloc7.Job = alloc.Job
	alloc7.JobID = alloc.JobID

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1003, []*structs.Allocation{alloc7}); err != nil {
		t.Fatalf("err: %v", err)
	}
	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
	if summary.Summary["web"].Starting != 1 || summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
		t.Fatalf("bad job summary: %v", summary)
	}
}

// Test that nonexistent deployment can't be updated
func TestStateStore_UpsertDeploymentStatusUpdate_Nonexistent(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Update the nonexistent deployment
	req := &structs.DeploymentStatusUpdateRequest{
		DeploymentUpdate: &structs.DeploymentStatusUpdate{
			DeploymentID: uuid.Generate(),
			Status:       structs.DeploymentStatusRunning,
		},
	}
	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "does not exist") {
		t.Fatalf("expected error updating the status because the deployment doesn't exist")
	}
}

// Test that terminal deployment can't be updated
func TestStateStore_UpsertDeploymentStatusUpdate_Terminal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a terminal deployment
	d := mock.Deployment()
	d.Status = structs.DeploymentStatusFailed

	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Update the deployment
	req := &structs.DeploymentStatusUpdateRequest{
		DeploymentUpdate: &structs.DeploymentStatusUpdate{
			DeploymentID: d.ID,
			Status:       structs.DeploymentStatusRunning,
		},
	}
	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
		t.Fatalf("expected error updating the status because the deployment is terminal")
	}
}

// Test that a non terminal deployment is updated and that a job and eval are
// created.
func TestStateStore_UpsertDeploymentStatusUpdate_NonTerminal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a deployment
	d := mock.Deployment()
	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create an eval and a job
	e := mock.Eval()
	j := mock.Job()

	// Update the deployment
	status, desc := structs.DeploymentStatusFailed, "foo"
	req := &structs.DeploymentStatusUpdateRequest{
		DeploymentUpdate: &structs.DeploymentStatusUpdate{
			DeploymentID:      d.ID,
			Status:            status,
			StatusDescription: desc,
		},
		Job:  j,
		Eval: e,
	}
	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Check that the status was updated properly
	ws := memdb.NewWatchSet()
	dout, err := state.DeploymentByID(ws, d.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if dout.Status != status || dout.StatusDescription != desc {
		t.Fatalf("bad: %#v", dout)
	}

	// Check that the evaluation was created
	eout, _ := state.EvalByID(ws, e.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if eout == nil {
		t.Fatalf("bad: %#v", eout)
	}

	// Check that the job was created
	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if jout == nil {
		t.Fatalf("bad: %#v", jout)
	}
}

// Test that when a deployment is updated to successful the job is updated to
// stable
func TestStateStore_UpsertDeploymentStatusUpdate_Successful(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a job
	job := mock.Job()
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1, job); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Insert a deployment
	d := structs.NewDeployment(job)
	if err := state.UpsertDeployment(2, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Update the deployment
	req := &structs.DeploymentStatusUpdateRequest{
		DeploymentUpdate: &structs.DeploymentStatusUpdate{
			DeploymentID:      d.ID,
			Status:            structs.DeploymentStatusSuccessful,
			StatusDescription: structs.DeploymentStatusDescriptionSuccessful,
		},
	}
	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 3, req)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Check that the status was updated properly
	ws := memdb.NewWatchSet()
	dout, err := state.DeploymentByID(ws, d.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if dout.Status != structs.DeploymentStatusSuccessful ||
		dout.StatusDescription != structs.DeploymentStatusDescriptionSuccessful {
		t.Fatalf("bad: %#v", dout)
	}

	// Check that the job was created
	jout, _ := state.JobByID(ws, job.Namespace, job.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if jout == nil {
		t.Fatalf("bad: %#v", jout)
	}
	if !jout.Stable {
		t.Fatalf("job not marked stable %#v", jout)
	}
	if jout.Version != d.JobVersion {
		t.Fatalf("job version changed; got %d; want %d", jout.Version, d.JobVersion)
	}
}

func TestStateStore_UpdateJobStability(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a job twice to get two versions
	job := mock.Job()
	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 1, job))

	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 2, job.Copy()))

	// Update the stability to true
	err := state.UpdateJobStability(3, job.Namespace, job.ID, 0, true)
	require.NoError(t, err)

	// Check that the job was updated properly
	ws := memdb.NewWatchSet()
	jout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
	require.NoError(t, err)
	require.NotNil(t, jout)
	require.True(t, jout.Stable, "job not marked as stable")

	// Update the stability to false
	err = state.UpdateJobStability(3, job.Namespace, job.ID, 0, false)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Check that the job was updated properly
	jout, err = state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
	require.NoError(t, err)
	require.NotNil(t, jout)
	require.False(t, jout.Stable)
}

// Test that nonexistent deployment can't be promoted
func TestStateStore_UpsertDeploymentPromotion_Nonexistent(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Promote the nonexistent deployment
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: uuid.Generate(),
			All:          true,
		},
	}
	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "does not exist") {
		t.Fatalf("expected error promoting because the deployment doesn't exist")
	}
}

// Test that terminal deployment can't be updated
func TestStateStore_UpsertDeploymentPromotion_Terminal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a terminal deployment
	d := mock.Deployment()
	d.Status = structs.DeploymentStatusFailed

	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Promote the deployment
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: d.ID,
			All:          true,
		},
	}
	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
		t.Fatalf("expected error updating the status because the deployment is terminal: %v", err)
	}
}

// Test promoting unhealthy canaries in a deployment.
func TestStateStore_UpsertDeploymentPromotion_Unhealthy(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	require := require.New(t)

	// Create a job
	j := mock.Job()
	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))

	// Create a deployment
	d := mock.Deployment()
	d.JobID = j.ID
	d.TaskGroups["web"].DesiredCanaries = 2
	require.Nil(state.UpsertDeployment(2, d))

	// Create a set of allocations
	c1 := mock.Alloc()
	c1.JobID = j.ID
	c1.DeploymentID = d.ID
	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
	c2 := mock.Alloc()
	c2.JobID = j.ID
	c2.DeploymentID = d.ID
	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)

	// Create a healthy but terminal alloc
	c3 := mock.Alloc()
	c3.JobID = j.ID
	c3.DeploymentID = d.ID
	c3.DesiredStatus = structs.AllocDesiredStatusStop
	c3.DeploymentStatus = &structs.AllocDeploymentStatus{Healthy: helper.BoolToPtr(true)}
	d.TaskGroups[c3.TaskGroup].PlacedCanaries = append(d.TaskGroups[c3.TaskGroup].PlacedCanaries, c3.ID)

	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2, c3}))

	// Promote the canaries
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: d.ID,
			All:          true,
		},
	}
	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
	require.NotNil(err)
	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
}

// Test promoting a deployment with no canaries
func TestStateStore_UpsertDeploymentPromotion_NoCanaries(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	require := require.New(t)

	// Create a job
	j := mock.Job()
	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))

	// Create a deployment
	d := mock.Deployment()
	d.TaskGroups["web"].DesiredCanaries = 2
	d.JobID = j.ID
	require.Nil(state.UpsertDeployment(2, d))

	// Promote the canaries
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: d.ID,
			All:          true,
		},
	}
	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
	require.NotNil(err)
	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
}

// Test promoting all canaries in a deployment.
func TestStateStore_UpsertDeploymentPromotion_All(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create a job with two task groups
	j := mock.Job()
	tg1 := j.TaskGroups[0]
	tg2 := tg1.Copy()
	tg2.Name = "foo"
	j.TaskGroups = append(j.TaskGroups, tg2)
	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1, j); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create a deployment
	d := mock.Deployment()
	d.StatusDescription = structs.DeploymentStatusDescriptionRunningNeedsPromotion
	d.JobID = j.ID
	d.TaskGroups = map[string]*structs.DeploymentState{
		"web": {
			DesiredTotal:    10,
			DesiredCanaries: 1,
		},
		"foo": {
			DesiredTotal:    10,
			DesiredCanaries: 1,
		},
	}
	if err := state.UpsertDeployment(2, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create a set of allocations
	c1 := mock.Alloc()
	c1.JobID = j.ID
	c1.DeploymentID = d.ID
	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(true),
	}
	c2 := mock.Alloc()
	c2.JobID = j.ID
	c2.DeploymentID = d.ID
	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
	c2.TaskGroup = tg2.Name
	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(true),
	}

	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create an eval
	e := mock.Eval()

	// Promote the canaries
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: d.ID,
			All:          true,
		},
		Eval: e,
	}
	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Check that the status per task group was updated properly
	ws := memdb.NewWatchSet()
	dout, err := state.DeploymentByID(ws, d.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if dout.StatusDescription != structs.DeploymentStatusDescriptionRunning {
		t.Fatalf("status description not updated: got %v; want %v", dout.StatusDescription, structs.DeploymentStatusDescriptionRunning)
	}
	if len(dout.TaskGroups) != 2 {
		t.Fatalf("bad: %#v", dout.TaskGroups)
	}
	for tg, state := range dout.TaskGroups {
		if !state.Promoted {
			t.Fatalf("bad: group %q not promoted %#v", tg, state)
		}
	}

	// Check that the evaluation was created
	eout, _ := state.EvalByID(ws, e.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if eout == nil {
		t.Fatalf("bad: %#v", eout)
	}
}

// Test promoting a subset of canaries in a deployment.
func TestStateStore_UpsertDeploymentPromotion_Subset(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)

	// Create a job with two task groups
	j := mock.Job()
	tg1 := j.TaskGroups[0]
	tg2 := tg1.Copy()
	tg2.Name = "foo"
	j.TaskGroups = append(j.TaskGroups, tg2)
	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))

	// Create a deployment
	d := mock.Deployment()
	d.JobID = j.ID
	d.TaskGroups = map[string]*structs.DeploymentState{
		"web": {
			DesiredTotal:    10,
			DesiredCanaries: 1,
		},
		"foo": {
			DesiredTotal:    10,
			DesiredCanaries: 1,
		},
	}
	require.Nil(state.UpsertDeployment(2, d))

	// Create a set of allocations for both groups, including an unhealthy one
	c1 := mock.Alloc()
	c1.JobID = j.ID
	c1.DeploymentID = d.ID
	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(true),
		Canary:  true,
	}

	// Should still be a canary
	c2 := mock.Alloc()
	c2.JobID = j.ID
	c2.DeploymentID = d.ID
	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
	c2.TaskGroup = tg2.Name
	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(true),
		Canary:  true,
	}

	c3 := mock.Alloc()
	c3.JobID = j.ID
	c3.DeploymentID = d.ID
	d.TaskGroups[c3.TaskGroup].PlacedCanaries = append(d.TaskGroups[c3.TaskGroup].PlacedCanaries, c3.ID)
	c3.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(false),
		Canary:  true,
	}

	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2, c3}))

	// Create an eval
	e := mock.Eval()

	// Promote the canaries
	req := &structs.ApplyDeploymentPromoteRequest{
		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
			DeploymentID: d.ID,
			Groups:       []string{"web"},
		},
		Eval: e,
	}
	require.Nil(state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req))

	// Check that the status per task group was updated properly
	ws := memdb.NewWatchSet()
	dout, err := state.DeploymentByID(ws, d.ID)
	require.Nil(err)
	require.Len(dout.TaskGroups, 2)
	require.Contains(dout.TaskGroups, "web")
	require.True(dout.TaskGroups["web"].Promoted)

	// Check that the evaluation was created
	eout, err := state.EvalByID(ws, e.ID)
	require.Nil(err)
	require.NotNil(eout)

	// Check the canary field was set properly
	aout1, err1 := state.AllocByID(ws, c1.ID)
	aout2, err2 := state.AllocByID(ws, c2.ID)
	aout3, err3 := state.AllocByID(ws, c3.ID)
	require.Nil(err1)
	require.Nil(err2)
	require.Nil(err3)
	require.NotNil(aout1)
	require.NotNil(aout2)
	require.NotNil(aout3)
	require.False(aout1.DeploymentStatus.Canary)
	require.True(aout2.DeploymentStatus.Canary)
	require.True(aout3.DeploymentStatus.Canary)
}

// Test that allocation health can't be set against a nonexistent deployment
func TestStateStore_UpsertDeploymentAllocHealth_Nonexistent(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Set health against the nonexistent deployment
	req := &structs.ApplyDeploymentAllocHealthRequest{
		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
			DeploymentID:         uuid.Generate(),
			HealthyAllocationIDs: []string{uuid.Generate()},
		},
	}
	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "does not exist") {
		t.Fatalf("expected error because the deployment doesn't exist: %v", err)
	}
}

// Test that allocation health can't be set against a terminal deployment
func TestStateStore_UpsertDeploymentAllocHealth_Terminal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a terminal deployment
	d := mock.Deployment()
	d.Status = structs.DeploymentStatusFailed

	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Set health against the terminal deployment
	req := &structs.ApplyDeploymentAllocHealthRequest{
		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
			DeploymentID:         d.ID,
			HealthyAllocationIDs: []string{uuid.Generate()},
		},
	}
	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
		t.Fatalf("expected error because the deployment is terminal: %v", err)
	}
}

// Test that allocation health can't be set against a nonexistent alloc
func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_Nonexistent(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a deployment
	d := mock.Deployment()
	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Set health against the terminal deployment
	req := &structs.ApplyDeploymentAllocHealthRequest{
		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
			DeploymentID:         d.ID,
			HealthyAllocationIDs: []string{uuid.Generate()},
		},
	}
	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
	if err == nil || !strings.Contains(err.Error(), "unknown alloc") {
		t.Fatalf("expected error because the alloc doesn't exist: %v", err)
	}
}

// Test that a deployments PlacedCanaries is properly updated
func TestStateStore_UpsertDeploymentAlloc_Canaries(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create a deployment
	d1 := mock.Deployment()
	require.NoError(t, state.UpsertDeployment(2, d1))

	// Create a Job
	job := mock.Job()
	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 3, job))

	// Create alloc with canary status
	a := mock.Alloc()
	a.JobID = job.ID
	a.DeploymentID = d1.ID
	a.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(false),
		Canary:  true,
	}
	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{a}))

	// Pull the deployment from state
	ws := memdb.NewWatchSet()
	deploy, err := state.DeploymentByID(ws, d1.ID)
	require.NoError(t, err)

	// Ensure that PlacedCanaries is accurate
	require.Equal(t, 1, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))

	// Create alloc without canary status
	b := mock.Alloc()
	b.JobID = job.ID
	b.DeploymentID = d1.ID
	b.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(false),
		Canary:  false,
	}
	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{b}))

	// Pull the deployment from state
	ws = memdb.NewWatchSet()
	deploy, err = state.DeploymentByID(ws, d1.ID)
	require.NoError(t, err)

	// Ensure that PlacedCanaries is accurate
	require.Equal(t, 1, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))

	// Create a second deployment
	d2 := mock.Deployment()
	require.NoError(t, state.UpsertDeployment(5, d2))

	c := mock.Alloc()
	c.JobID = job.ID
	c.DeploymentID = d2.ID
	c.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(false),
		Canary:  true,
	}
	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 6, []*structs.Allocation{c}))

	ws = memdb.NewWatchSet()
	deploy2, err := state.DeploymentByID(ws, d2.ID)
	require.NoError(t, err)

	// Ensure that PlacedCanaries is accurate
	require.Equal(t, 1, len(deploy2.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
}

func TestStateStore_UpsertDeploymentAlloc_NoCanaries(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Create a deployment
	d1 := mock.Deployment()
	require.NoError(t, state.UpsertDeployment(2, d1))

	// Create a Job
	job := mock.Job()
	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 3, job))

	// Create alloc with canary status
	a := mock.Alloc()
	a.JobID = job.ID
	a.DeploymentID = d1.ID
	a.DeploymentStatus = &structs.AllocDeploymentStatus{
		Healthy: helper.BoolToPtr(true),
		Canary:  false,
	}
	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{a}))

	// Pull the deployment from state
	ws := memdb.NewWatchSet()
	deploy, err := state.DeploymentByID(ws, d1.ID)
	require.NoError(t, err)

	// Ensure that PlacedCanaries is accurate
	require.Equal(t, 0, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
}

// Test that allocation health can't be set for an alloc with mismatched
// deployment ids
func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_MismatchDeployment(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert two  deployment
	d1 := mock.Deployment()
	d2 := mock.Deployment()
	if err := state.UpsertDeployment(1, d1); err != nil {
		t.Fatalf("bad: %v", err)
	}
	if err := state.UpsertDeployment(2, d2); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Insert an alloc for a random deployment
	a := mock.Alloc()
	a.DeploymentID = d1.ID
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{a}); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Set health against the terminal deployment
	req := &structs.ApplyDeploymentAllocHealthRequest{
		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
			DeploymentID:         d2.ID,
			HealthyAllocationIDs: []string{a.ID},
		},
	}
	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 4, req)
	if err == nil || !strings.Contains(err.Error(), "not part of deployment") {
		t.Fatalf("expected error because the alloc isn't part of the deployment: %v", err)
	}
}

// Test that allocation health is properly set
func TestStateStore_UpsertDeploymentAllocHealth(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)

	// Insert a deployment
	d := mock.Deployment()
	d.TaskGroups["web"].ProgressDeadline = 5 * time.Minute
	if err := state.UpsertDeployment(1, d); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Insert two allocations
	a1 := mock.Alloc()
	a1.DeploymentID = d.ID
	a2 := mock.Alloc()
	a2.DeploymentID = d.ID
	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 2, []*structs.Allocation{a1, a2}); err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Create a job to roll back to
	j := mock.Job()

	// Create an eval that should be upserted
	e := mock.Eval()

	// Create a status update for the deployment
	status, desc := structs.DeploymentStatusFailed, "foo"
	u := &structs.DeploymentStatusUpdate{
		DeploymentID:      d.ID,
		Status:            status,
		StatusDescription: desc,
	}

	// Capture the time for the update
	ts := time.Now()

	// Set health against the deployment
	req := &structs.ApplyDeploymentAllocHealthRequest{
		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
			DeploymentID:           d.ID,
			HealthyAllocationIDs:   []string{a1.ID},
			UnhealthyAllocationIDs: []string{a2.ID},
		},
		Job:              j,
		Eval:             e,
		DeploymentUpdate: u,
		Timestamp:        ts,
	}
	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 3, req)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}

	// Check that the status was updated properly
	ws := memdb.NewWatchSet()
	dout, err := state.DeploymentByID(ws, d.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if dout.Status != status || dout.StatusDescription != desc {
		t.Fatalf("bad: %#v", dout)
	}

	// Check that the evaluation was created
	eout, _ := state.EvalByID(ws, e.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if eout == nil {
		t.Fatalf("bad: %#v", eout)
	}

	// Check that the job was created
	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
	if err != nil {
		t.Fatalf("bad: %v", err)
	}
	if jout == nil {
		t.Fatalf("bad: %#v", jout)
	}

	// Check the status of the allocs
	out1, err := state.AllocByID(ws, a1.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	out2, err := state.AllocByID(ws, a2.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !out1.DeploymentStatus.IsHealthy() {
		t.Fatalf("bad: alloc %q not healthy", out1.ID)
	}
	if !out2.DeploymentStatus.IsUnhealthy() {
		t.Fatalf("bad: alloc %q not unhealthy", out2.ID)
	}

	if !out1.DeploymentStatus.Timestamp.Equal(ts) {
		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out1.ID, out1.DeploymentStatus.Timestamp, ts)
	}
	if !out2.DeploymentStatus.Timestamp.Equal(ts) {
		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out2.ID, out2.DeploymentStatus.Timestamp, ts)
	}
}

func TestStateStore_UpsertVaultAccessors(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	a := mock.VaultAccessor()
	a2 := mock.VaultAccessor()

	ws := memdb.NewWatchSet()
	if _, err := state.VaultAccessor(ws, a.Accessor); err != nil {
		t.Fatalf("err: %v", err)
	}

	if _, err := state.VaultAccessor(ws, a2.Accessor); err != nil {
		t.Fatalf("err: %v", err)
	}

	err := state.UpsertVaultAccessor(1000, []*structs.VaultAccessor{a, a2})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.VaultAccessor(ws, a.Accessor)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(a, out) {
		t.Fatalf("bad: %#v %#v", a, out)
	}

	out, err = state.VaultAccessor(ws, a2.Accessor)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(a2, out) {
		t.Fatalf("bad: %#v %#v", a2, out)
	}

	iter, err := state.VaultAccessors(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}

		count++
		accessor := raw.(*structs.VaultAccessor)

		if !reflect.DeepEqual(accessor, a) && !reflect.DeepEqual(accessor, a2) {
			t.Fatalf("bad: %#v", accessor)
		}
	}

	if count != 2 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("vault_accessors")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteVaultAccessors(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	a1 := mock.VaultAccessor()
	a2 := mock.VaultAccessor()
	accessors := []*structs.VaultAccessor{a1, a2}

	err := state.UpsertVaultAccessor(1000, accessors)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	if _, err := state.VaultAccessor(ws, a1.Accessor); err != nil {
		t.Fatalf("err: %v", err)
	}

	err = state.DeleteVaultAccessors(1001, accessors)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.VaultAccessor(ws, a1.Accessor)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if out != nil {
		t.Fatalf("bad: %#v %#v", a1, out)
	}
	out, err = state.VaultAccessor(ws, a2.Accessor)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if out != nil {
		t.Fatalf("bad: %#v %#v", a2, out)
	}

	index, err := state.Index("vault_accessors")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_VaultAccessorsByAlloc(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	var accessors []*structs.VaultAccessor
	var expected []*structs.VaultAccessor

	for i := 0; i < 5; i++ {
		accessor := mock.VaultAccessor()
		accessor.AllocID = alloc.ID
		expected = append(expected, accessor)
		accessors = append(accessors, accessor)
	}

	for i := 0; i < 10; i++ {
		accessor := mock.VaultAccessor()
		accessors = append(accessors, accessor)
	}

	err := state.UpsertVaultAccessor(1000, accessors)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.VaultAccessorsByAlloc(ws, alloc.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if len(expected) != len(out) {
		t.Fatalf("bad: %#v %#v", len(expected), len(out))
	}

	index, err := state.Index("vault_accessors")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_VaultAccessorsByNode(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	node := mock.Node()
	var accessors []*structs.VaultAccessor
	var expected []*structs.VaultAccessor

	for i := 0; i < 5; i++ {
		accessor := mock.VaultAccessor()
		accessor.NodeID = node.ID
		expected = append(expected, accessor)
		accessors = append(accessors, accessor)
	}

	for i := 0; i < 10; i++ {
		accessor := mock.VaultAccessor()
		accessors = append(accessors, accessor)
	}

	err := state.UpsertVaultAccessor(1000, accessors)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	ws := memdb.NewWatchSet()
	out, err := state.VaultAccessorsByNode(ws, node.ID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if len(expected) != len(out) {
		t.Fatalf("bad: %#v %#v", len(expected), len(out))
	}

	index, err := state.Index("vault_accessors")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_RestoreVaultAccessor(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	a := mock.VaultAccessor()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.VaultAccessorRestore(a)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.VaultAccessor(ws, a.Accessor)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if !reflect.DeepEqual(out, a) {
		t.Fatalf("Bad: %#v %#v", out, a)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_UpsertSITokenAccessors(t *testing.T) {
	t.Parallel()
	r := require.New(t)

	state := testStateStore(t)
	a1 := mock.SITokenAccessor()
	a2 := mock.SITokenAccessor()

	ws := memdb.NewWatchSet()
	var err error

	_, err = state.SITokenAccessor(ws, a1.AccessorID)
	r.NoError(err)

	_, err = state.SITokenAccessor(ws, a2.AccessorID)
	r.NoError(err)

	err = state.UpsertSITokenAccessors(1000, []*structs.SITokenAccessor{a1, a2})
	r.NoError(err)

	wsFired := watchFired(ws)
	r.True(wsFired)

	noInsertWS := memdb.NewWatchSet()
	result1, err := state.SITokenAccessor(noInsertWS, a1.AccessorID)
	r.NoError(err)
	r.Equal(a1, result1)

	result2, err := state.SITokenAccessor(noInsertWS, a2.AccessorID)
	r.NoError(err)
	r.Equal(a2, result2)

	iter, err := state.SITokenAccessors(noInsertWS)
	r.NoError(err)

	count := 0
	for raw := iter.Next(); raw != nil; raw = iter.Next() {
		count++
		accessor := raw.(*structs.SITokenAccessor)
		// iterator is sorted by dynamic UUID
		matches := reflect.DeepEqual(a1, accessor) || reflect.DeepEqual(a2, accessor)
		r.True(matches)
	}
	r.Equal(2, count)

	index, err := state.Index(siTokenAccessorTable)
	r.NoError(err)
	r.Equal(uint64(1000), index)

	noInsertWSFired := watchFired(noInsertWS)
	r.False(noInsertWSFired)
}

func TestStateStore_DeleteSITokenAccessors(t *testing.T) {
	t.Parallel()
	r := require.New(t)

	state := testStateStore(t)
	a1 := mock.SITokenAccessor()
	a2 := mock.SITokenAccessor()
	accessors := []*structs.SITokenAccessor{a1, a2}
	var err error

	err = state.UpsertSITokenAccessors(1000, accessors)
	r.NoError(err)

	ws := memdb.NewWatchSet()
	_, err = state.SITokenAccessor(ws, a1.AccessorID)
	r.NoError(err)

	err = state.DeleteSITokenAccessors(1001, accessors)
	r.NoError(err)

	wsFired := watchFired(ws)
	r.True(wsFired)

	wsPostDelete := memdb.NewWatchSet()

	result1, err := state.SITokenAccessor(wsPostDelete, a1.AccessorID)
	r.NoError(err)
	r.Nil(result1) // was deleted

	result2, err := state.SITokenAccessor(wsPostDelete, a2.AccessorID)
	r.NoError(err)
	r.Nil(result2) // was deleted

	index, err := state.Index(siTokenAccessorTable)
	r.NoError(err)
	r.Equal(uint64(1001), index)

	wsPostDeleteFired := watchFired(wsPostDelete)
	r.False(wsPostDeleteFired)
}

func TestStateStore_SITokenAccessorsByAlloc(t *testing.T) {
	t.Parallel()
	r := require.New(t)

	state := testStateStore(t)
	alloc := mock.Alloc()
	var accessors []*structs.SITokenAccessor
	var expected []*structs.SITokenAccessor

	for i := 0; i < 5; i++ {
		accessor := mock.SITokenAccessor()
		accessor.AllocID = alloc.ID
		expected = append(expected, accessor)
		accessors = append(accessors, accessor)
	}

	for i := 0; i < 10; i++ {
		accessor := mock.SITokenAccessor()
		accessor.AllocID = uuid.Generate() // does not belong to alloc
		accessors = append(accessors, accessor)
	}

	err := state.UpsertSITokenAccessors(1000, accessors)
	r.NoError(err)

	ws := memdb.NewWatchSet()
	result, err := state.SITokenAccessorsByAlloc(ws, alloc.ID)
	r.NoError(err)
	r.ElementsMatch(expected, result)

	index, err := state.Index(siTokenAccessorTable)
	r.NoError(err)
	r.Equal(uint64(1000), index)

	wsFired := watchFired(ws)
	r.False(wsFired)
}

func TestStateStore_SITokenAccessorsByNode(t *testing.T) {
	t.Parallel()
	r := require.New(t)

	state := testStateStore(t)
	node := mock.Node()
	var accessors []*structs.SITokenAccessor
	var expected []*structs.SITokenAccessor
	var err error

	for i := 0; i < 5; i++ {
		accessor := mock.SITokenAccessor()
		accessor.NodeID = node.ID
		expected = append(expected, accessor)
		accessors = append(accessors, accessor)
	}

	for i := 0; i < 10; i++ {
		accessor := mock.SITokenAccessor()
		accessor.NodeID = uuid.Generate() // does not belong to node
		accessors = append(accessors, accessor)
	}

	err = state.UpsertSITokenAccessors(1000, accessors)
	r.NoError(err)

	ws := memdb.NewWatchSet()
	result, err := state.SITokenAccessorsByNode(ws, node.ID)
	r.NoError(err)
	r.ElementsMatch(expected, result)

	index, err := state.Index(siTokenAccessorTable)
	r.NoError(err)
	r.Equal(uint64(1000), index)

	wsFired := watchFired(ws)
	r.False(wsFired)
}

func TestStateStore_RestoreSITokenAccessor(t *testing.T) {
	t.Parallel()
	r := require.New(t)

	state := testStateStore(t)
	a1 := mock.SITokenAccessor()

	restore, err := state.Restore()
	r.NoError(err)

	err = restore.SITokenAccessorRestore(a1)
	r.NoError(err)

	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	result, err := state.SITokenAccessor(ws, a1.AccessorID)
	r.NoError(err)
	r.Equal(a1, result)

	wsFired := watchFired(ws)
	r.False(wsFired)
}

func TestStateStore_UpsertACLPolicy(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	policy := mock.ACLPolicy()
	policy2 := mock.ACLPolicy()

	ws := memdb.NewWatchSet()
	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
		t.Fatalf("err: %v", err)
	}
	if _, err := state.ACLPolicyByName(ws, policy2.Name); err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy, policy2}); err != nil {
		t.Fatalf("err: %v", err)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.ACLPolicyByName(ws, policy.Name)
	assert.Equal(t, nil, err)
	assert.Equal(t, policy, out)

	out, err = state.ACLPolicyByName(ws, policy2.Name)
	assert.Equal(t, nil, err)
	assert.Equal(t, policy2, out)

	iter, err := state.ACLPolicies(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both policies
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 2 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("acl_policy")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteACLPolicy(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	policy := mock.ACLPolicy()
	policy2 := mock.ACLPolicy()

	// Create the policy
	if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy, policy2}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watcher
	ws := memdb.NewWatchSet()
	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Delete the policy
	if err := state.DeleteACLPolicies(structs.MsgTypeTestSetup, 1001, []string{policy.Name, policy2.Name}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure watching triggered
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	// Ensure we don't get the object back
	ws = memdb.NewWatchSet()
	out, err := state.ACLPolicyByName(ws, policy.Name)
	assert.Equal(t, nil, err)
	if out != nil {
		t.Fatalf("bad: %#v", out)
	}

	iter, err := state.ACLPolicies(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see neither policy
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 0 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("acl_policy")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_ACLPolicyByNamePrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	names := []string{
		"foo",
		"bar",
		"foobar",
		"foozip",
		"zip",
	}

	// Create the policies
	var baseIndex uint64 = 1000
	for _, name := range names {
		p := mock.ACLPolicy()
		p.Name = name
		if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, baseIndex, []*structs.ACLPolicy{p}); err != nil {
			t.Fatalf("err: %v", err)
		}
		baseIndex++
	}

	// Scan by prefix
	iter, err := state.ACLPolicyByNamePrefix(nil, "foo")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both policies
	count := 0
	out := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
		out = append(out, raw.(*structs.ACLPolicy).Name)
	}
	if count != 3 {
		t.Fatalf("bad: %d %v", count, out)
	}
	sort.Strings(out)

	expect := []string{"foo", "foobar", "foozip"}
	assert.Equal(t, expect, out)
}

func TestStateStore_BootstrapACLTokens(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	tk1 := mock.ACLToken()
	tk2 := mock.ACLToken()

	ok, resetIdx, err := state.CanBootstrapACLToken()
	assert.Nil(t, err)
	assert.Equal(t, true, ok)
	assert.EqualValues(t, 0, resetIdx)

	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1000, 0, tk1); err != nil {
		t.Fatalf("err: %v", err)
	}

	out, err := state.ACLTokenByAccessorID(nil, tk1.AccessorID)
	assert.Equal(t, nil, err)
	assert.Equal(t, tk1, out)

	ok, resetIdx, err = state.CanBootstrapACLToken()
	assert.Nil(t, err)
	assert.Equal(t, false, ok)
	assert.EqualValues(t, 1000, resetIdx)

	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1001, 0, tk2); err == nil {
		t.Fatalf("expected error")
	}

	iter, err := state.ACLTokens(nil)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both policies
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 1 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("acl_token")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}
	index, err = state.Index("acl_token_bootstrap")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	// Should allow bootstrap with reset index
	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1001, 1000, tk2); err != nil {
		t.Fatalf("err %v", err)
	}

	// Check we've modified the index
	index, err = state.Index("acl_token")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}
	index, err = state.Index("acl_token_bootstrap")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}
}

func TestStateStore_UpsertACLTokens(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	tk1 := mock.ACLToken()
	tk2 := mock.ACLToken()

	ws := memdb.NewWatchSet()
	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
		t.Fatalf("err: %v", err)
	}
	if _, err := state.ACLTokenByAccessorID(ws, tk2.AccessorID); err != nil {
		t.Fatalf("err: %v", err)
	}

	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2}); err != nil {
		t.Fatalf("err: %v", err)
	}
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	ws = memdb.NewWatchSet()
	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
	assert.Equal(t, nil, err)
	assert.Equal(t, tk1, out)

	out, err = state.ACLTokenByAccessorID(ws, tk2.AccessorID)
	assert.Equal(t, nil, err)
	assert.Equal(t, tk2, out)

	out, err = state.ACLTokenBySecretID(ws, tk1.SecretID)
	assert.Equal(t, nil, err)
	assert.Equal(t, tk1, out)

	out, err = state.ACLTokenBySecretID(ws, tk2.SecretID)
	assert.Equal(t, nil, err)
	assert.Equal(t, tk2, out)

	iter, err := state.ACLTokens(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both policies
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 2 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("acl_token")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1000 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_DeleteACLTokens(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	tk1 := mock.ACLToken()
	tk2 := mock.ACLToken()

	// Create the tokens
	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Create a watcher
	ws := memdb.NewWatchSet()
	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Delete the token
	if err := state.DeleteACLTokens(structs.MsgTypeTestSetup, 1001, []string{tk1.AccessorID, tk2.AccessorID}); err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure watching triggered
	if !watchFired(ws) {
		t.Fatalf("bad")
	}

	// Ensure we don't get the object back
	ws = memdb.NewWatchSet()
	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
	assert.Equal(t, nil, err)
	if out != nil {
		t.Fatalf("bad: %#v", out)
	}

	iter, err := state.ACLTokens(ws)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both policies
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 0 {
		t.Fatalf("bad: %d", count)
	}

	index, err := state.Index("acl_token")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if index != 1001 {
		t.Fatalf("bad: %d", index)
	}

	if watchFired(ws) {
		t.Fatalf("bad")
	}
}

func TestStateStore_ACLTokenByAccessorIDPrefix(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	prefixes := []string{
		"aaaa",
		"aabb",
		"bbbb",
		"bbcc",
		"ffff",
	}

	// Create the tokens
	var baseIndex uint64 = 1000
	for _, prefix := range prefixes {
		tk := mock.ACLToken()
		tk.AccessorID = prefix + tk.AccessorID[4:]
		if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, baseIndex, []*structs.ACLToken{tk}); err != nil {
			t.Fatalf("err: %v", err)
		}
		baseIndex++
	}

	// Scan by prefix
	iter, err := state.ACLTokenByAccessorIDPrefix(nil, "aa")
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see both tokens
	count := 0
	out := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
		out = append(out, raw.(*structs.ACLToken).AccessorID[:4])
	}
	if count != 2 {
		t.Fatalf("bad: %d %v", count, out)
	}
	sort.Strings(out)

	expect := []string{"aaaa", "aabb"}
	assert.Equal(t, expect, out)
}

func TestStateStore_RestoreACLPolicy(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	policy := mock.ACLPolicy()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.ACLPolicyRestore(policy)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.ACLPolicyByName(ws, policy.Name)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	assert.Equal(t, policy, out)
}

func TestStateStore_ACLTokensByGlobal(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	tk1 := mock.ACLToken()
	tk2 := mock.ACLToken()
	tk3 := mock.ACLToken()
	tk4 := mock.ACLToken()
	tk3.Global = true

	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2, tk3, tk4}); err != nil {
		t.Fatalf("err: %v", err)
	}

	iter, err := state.ACLTokensByGlobal(nil, true)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	// Ensure we see the one global policies
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	if count != 1 {
		t.Fatalf("bad: %d", count)
	}
}

func TestStateStore_RestoreACLToken(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	token := mock.ACLToken()

	restore, err := state.Restore()
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	err = restore.ACLTokenRestore(token)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	require.NoError(t, restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.ACLTokenByAccessorID(ws, token.AccessorID)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	assert.Equal(t, token, out)
}

func TestStateStore_OneTimeTokens(t *testing.T) {
	t.Parallel()
	index := uint64(100)
	state := testStateStore(t)

	// create some ACL tokens

	token1 := mock.ACLToken()
	token2 := mock.ACLToken()
	token3 := mock.ACLToken()
	index++
	require.Nil(t, state.UpsertACLTokens(
		structs.MsgTypeTestSetup, index,
		[]*structs.ACLToken{token1, token2, token3}))

	otts := []*structs.OneTimeToken{
		{
			// expired OTT for token1
			OneTimeSecretID: uuid.Generate(),
			AccessorID:      token1.AccessorID,
			ExpiresAt:       time.Now().Add(-1 * time.Minute),
		},
		{
			// valid OTT for token2
			OneTimeSecretID: uuid.Generate(),
			AccessorID:      token2.AccessorID,
			ExpiresAt:       time.Now().Add(10 * time.Minute),
		},
		{
			// new but expired OTT for token2; this will be accepted even
			// though it's expired and overwrite the other one
			OneTimeSecretID: uuid.Generate(),
			AccessorID:      token2.AccessorID,
			ExpiresAt:       time.Now().Add(-10 * time.Minute),
		},
		{
			// valid OTT for token3
			AccessorID:      token3.AccessorID,
			OneTimeSecretID: uuid.Generate(),
			ExpiresAt:       time.Now().Add(10 * time.Minute),
		},
		{
			// new valid OTT for token3
			OneTimeSecretID: uuid.Generate(),
			AccessorID:      token3.AccessorID,
			ExpiresAt:       time.Now().Add(5 * time.Minute),
		},
	}

	for _, ott := range otts {
		index++
		require.NoError(t, state.UpsertOneTimeToken(structs.MsgTypeTestSetup, index, ott))
	}

	// verify that we have exactly one OTT for each AccessorID

	txn := state.db.ReadTxn()
	iter, err := txn.Get("one_time_token", "id")
	require.NoError(t, err)
	results := []*structs.OneTimeToken{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		ott, ok := raw.(*structs.OneTimeToken)
		require.True(t, ok)
		results = append(results, ott)
	}

	// results aren't ordered but if we have 3 OTT and all 3 tokens, we know
	// we have no duplicate accessors
	require.Len(t, results, 3)
	accessors := []string{
		results[0].AccessorID, results[1].AccessorID, results[2].AccessorID}
	require.Contains(t, accessors, token1.AccessorID)
	require.Contains(t, accessors, token2.AccessorID)
	require.Contains(t, accessors, token3.AccessorID)

	// now verify expiration

	getExpiredTokens := func() []*structs.OneTimeToken {
		txn := state.db.ReadTxn()
		iter, err := state.oneTimeTokensExpiredTxn(txn, nil)
		require.NoError(t, err)

		results := []*structs.OneTimeToken{}
		for {
			raw := iter.Next()
			if raw == nil {
				break
			}
			ott, ok := raw.(*structs.OneTimeToken)
			require.True(t, ok)
			results = append(results, ott)
		}
		return results
	}

	results = getExpiredTokens()
	require.Len(t, results, 2)

	// results aren't ordered
	expiredAccessors := []string{results[0].AccessorID, results[1].AccessorID}
	require.Contains(t, expiredAccessors, token1.AccessorID)
	require.Contains(t, expiredAccessors, token2.AccessorID)
	require.True(t, time.Now().After(results[0].ExpiresAt))
	require.True(t, time.Now().After(results[1].ExpiresAt))

	// clear the expired tokens and verify they're gone
	index++
	require.NoError(t,
		state.ExpireOneTimeTokens(structs.MsgTypeTestSetup, index))

	results = getExpiredTokens()
	require.Len(t, results, 0)

	// query the unexpired token
	ott, err := state.OneTimeTokenBySecret(nil, otts[len(otts)-1].OneTimeSecretID)
	require.NoError(t, err)
	require.Equal(t, token3.AccessorID, ott.AccessorID)
	require.True(t, time.Now().Before(ott.ExpiresAt))

	restore, err := state.Restore()
	require.NoError(t, err)
	err = restore.OneTimeTokenRestore(ott)
	require.NoError(t, err)
	require.NoError(t, restore.Commit())

	ott, err = state.OneTimeTokenBySecret(nil, otts[len(otts)-1].OneTimeSecretID)
	require.NoError(t, err)
	require.Equal(t, token3.AccessorID, ott.AccessorID)
}

func TestStateStore_SchedulerConfig(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	schedConfig := &structs.SchedulerConfiguration{
		PreemptionConfig: structs.PreemptionConfig{
			SystemSchedulerEnabled: false,
		},
		CreateIndex: 100,
		ModifyIndex: 200,
	}

	require := require.New(t)
	restore, err := state.Restore()
	require.Nil(err)

	err = restore.SchedulerConfigRestore(schedConfig)
	require.Nil(err)

	require.NoError(restore.Commit())

	modIndex, out, err := state.SchedulerConfig()
	require.Nil(err)
	require.Equal(schedConfig.ModifyIndex, modIndex)

	require.Equal(schedConfig, out)
}

func TestStateStore_ClusterMetadata(t *testing.T) {
	require := require.New(t)

	state := testStateStore(t)
	clusterID := "12345678-1234-1234-1234-1234567890"
	now := time.Now().UnixNano()
	meta := &structs.ClusterMetadata{ClusterID: clusterID, CreateTime: now}

	err := state.ClusterSetMetadata(100, meta)
	require.NoError(err)

	result, err := state.ClusterMetadata(nil)
	require.NoError(err)
	require.Equal(clusterID, result.ClusterID)
	require.Equal(now, result.CreateTime)
}

func TestStateStore_ClusterMetadataRestore(t *testing.T) {
	require := require.New(t)

	state := testStateStore(t)
	clusterID := "12345678-1234-1234-1234-1234567890"
	now := time.Now().UnixNano()
	meta := &structs.ClusterMetadata{ClusterID: clusterID, CreateTime: now}

	restore, err := state.Restore()
	require.NoError(err)

	err = restore.ClusterMetadataRestore(meta)
	require.NoError(err)

	require.NoError(restore.Commit())

	out, err := state.ClusterMetadata(nil)
	require.NoError(err)
	require.Equal(clusterID, out.ClusterID)
	require.Equal(now, out.CreateTime)
}

func TestStateStore_RestoreScalingPolicy(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	scalingPolicy := mock.ScalingPolicy()

	restore, err := state.Restore()
	require.NoError(err)

	err = restore.ScalingPolicyRestore(scalingPolicy)
	require.NoError(err)
	require.NoError(restore.Commit())

	ws := memdb.NewWatchSet()
	out, err := state.ScalingPolicyByID(ws, scalingPolicy.ID)
	require.NoError(err)
	require.EqualValues(out, scalingPolicy)
}

func TestStateStore_UpsertScalingPolicy(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	policy := mock.ScalingPolicy()
	policy2 := mock.ScalingPolicy()

	wsAll := memdb.NewWatchSet()
	all, err := state.ScalingPolicies(wsAll)
	require.NoError(err)
	require.Nil(all.Next())

	ws := memdb.NewWatchSet()
	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.Nil(out)

	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
	require.NoError(err)
	require.Nil(out)

	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
	require.NoError(err)
	require.True(watchFired(ws))
	require.True(watchFired(wsAll))

	ws = memdb.NewWatchSet()
	out, err = state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.Equal(policy, out)

	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
	require.NoError(err)
	require.Equal(policy2, out)

	// Ensure we see both policies
	countPolicies := func() (n int, err error) {
		iter, err := state.ScalingPolicies(ws)
		if err != nil {
			return
		}

		for raw := iter.Next(); raw != nil; raw = iter.Next() {
			n++
		}
		return
	}

	count, err := countPolicies()
	require.NoError(err)
	require.Equal(2, count)

	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.True(1000 == index)
	require.False(watchFired(ws))

	// Check that we can add policy with same target but different type
	policy3 := mock.ScalingPolicy()
	for k, v := range policy2.Target {
		policy3.Target[k] = v
	}

	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy3})
	require.NoError(err)

	// Ensure we see both policies, since target didn't change
	count, err = countPolicies()
	require.NoError(err)
	require.Equal(2, count)

	// Change type and check if we see 3
	policy3.Type = "other-type"

	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy3})
	require.NoError(err)

	count, err = countPolicies()
	require.NoError(err)
	require.Equal(3, count)
}

func TestStateStore_UpsertScalingPolicy_Namespace(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	otherNamespace := "not-default-namespace"
	state := testStateStore(t)
	policy := mock.ScalingPolicy()
	policy2 := mock.ScalingPolicy()
	policy2.Target[structs.ScalingTargetNamespace] = otherNamespace

	ws1 := memdb.NewWatchSet()
	iter, err := state.ScalingPoliciesByNamespace(ws1, structs.DefaultNamespace, "")
	require.NoError(err)
	require.Nil(iter.Next())

	ws2 := memdb.NewWatchSet()
	iter, err = state.ScalingPoliciesByNamespace(ws2, otherNamespace, "")
	require.NoError(err)
	require.Nil(iter.Next())

	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
	require.NoError(err)
	require.True(watchFired(ws1))
	require.True(watchFired(ws2))

	iter, err = state.ScalingPoliciesByNamespace(nil, structs.DefaultNamespace, "")
	require.NoError(err)
	policiesInDefaultNamespace := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		policiesInDefaultNamespace = append(policiesInDefaultNamespace, raw.(*structs.ScalingPolicy).ID)
	}
	require.ElementsMatch([]string{policy.ID}, policiesInDefaultNamespace)

	iter, err = state.ScalingPoliciesByNamespace(nil, otherNamespace, "")
	require.NoError(err)
	policiesInOtherNamespace := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		policiesInOtherNamespace = append(policiesInOtherNamespace, raw.(*structs.ScalingPolicy).ID)
	}
	require.ElementsMatch([]string{policy2.ID}, policiesInOtherNamespace)
}

func TestStateStore_UpsertScalingPolicy_Namespace_PrefixBug(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	ns1 := "name"
	ns2 := "name2" // matches prefix "name"
	state := testStateStore(t)
	policy1 := mock.ScalingPolicy()
	policy1.Target[structs.ScalingTargetNamespace] = ns1
	policy2 := mock.ScalingPolicy()
	policy2.Target[structs.ScalingTargetNamespace] = ns2

	ws1 := memdb.NewWatchSet()
	iter, err := state.ScalingPoliciesByNamespace(ws1, ns1, "")
	require.NoError(err)
	require.Nil(iter.Next())

	ws2 := memdb.NewWatchSet()
	iter, err = state.ScalingPoliciesByNamespace(ws2, ns2, "")
	require.NoError(err)
	require.Nil(iter.Next())

	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy1, policy2})
	require.NoError(err)
	require.True(watchFired(ws1))
	require.True(watchFired(ws2))

	iter, err = state.ScalingPoliciesByNamespace(nil, ns1, "")
	require.NoError(err)
	policiesInNS1 := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		policiesInNS1 = append(policiesInNS1, raw.(*structs.ScalingPolicy).ID)
	}
	require.ElementsMatch([]string{policy1.ID}, policiesInNS1)

	iter, err = state.ScalingPoliciesByNamespace(nil, ns2, "")
	require.NoError(err)
	policiesInNS2 := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		policiesInNS2 = append(policiesInNS2, raw.(*structs.ScalingPolicy).ID)
	}
	require.ElementsMatch([]string{policy2.ID}, policiesInNS2)
}

// Scaling Policy IDs are generated randomly during Job.Register
// Subsequent updates of the job should preserve the ID for the scaling policy
// associated with a given target.
func TestStateStore_UpsertJob_PreserveScalingPolicyIDsAndIndex(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)
	job, policy := mock.JobWithScalingPolicy()

	var newIndex uint64 = 1000
	err := state.UpsertJob(structs.MsgTypeTestSetup, newIndex, job)
	require.NoError(err)

	ws := memdb.NewWatchSet()
	p1, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(p1)
	require.Equal(newIndex, p1.CreateIndex)
	require.Equal(newIndex, p1.ModifyIndex)

	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.Equal(newIndex, index)
	require.NotEmpty(p1.ID)

	// update the job
	job.Meta["new-meta"] = "new-value"
	newIndex += 100
	err = state.UpsertJob(structs.MsgTypeTestSetup, newIndex, job)
	require.NoError(err)
	require.False(watchFired(ws), "watch should not have fired")

	p2, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(p2)
	require.Equal(p1.ID, p2.ID, "ID should not have changed")
	require.Equal(p1.CreateIndex, p2.CreateIndex)
	require.Equal(p1.ModifyIndex, p2.ModifyIndex)

	index, err = state.Index("scaling_policy")
	require.NoError(err)
	require.Equal(index, p1.CreateIndex, "table index should not have changed")
}

// Updating the scaling policy for a job should update the index table and fire the watch.
// This test is the converse of TestStateStore_UpsertJob_PreserveScalingPolicyIDsAndIndex
func TestStateStore_UpsertJob_UpdateScalingPolicy(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)
	job, policy := mock.JobWithScalingPolicy()

	var oldIndex uint64 = 1000
	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, oldIndex, job))

	ws := memdb.NewWatchSet()
	p1, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(p1)
	require.Equal(oldIndex, p1.CreateIndex)
	require.Equal(oldIndex, p1.ModifyIndex)
	prevId := p1.ID

	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.Equal(oldIndex, index)
	require.NotEmpty(p1.ID)

	// update the job with the updated scaling policy; make sure to use a different object
	newPolicy := p1.Copy()
	newPolicy.Policy["new-field"] = "new-value"
	job.TaskGroups[0].Scaling = newPolicy
	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, oldIndex+100, job))
	require.True(watchFired(ws), "watch should have fired")

	p2, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(p2)
	require.Equal(p2.Policy["new-field"], "new-value")
	require.Equal(prevId, p2.ID, "ID should not have changed")
	require.Equal(oldIndex, p2.CreateIndex)
	require.Greater(p2.ModifyIndex, oldIndex, "ModifyIndex should have advanced")

	index, err = state.Index("scaling_policy")
	require.NoError(err)
	require.Greater(index, oldIndex, "table index should have advanced")
}

func TestStateStore_DeleteScalingPolicies(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)
	policy := mock.ScalingPolicy()
	policy2 := mock.ScalingPolicy()

	// Create the policy
	err := state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
	require.NoError(err)

	// Create a watcher
	ws := memdb.NewWatchSet()
	_, err = state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)

	// Delete the policy
	err = state.DeleteScalingPolicies(1001, []string{policy.ID, policy2.ID})
	require.NoError(err)

	// Ensure watching triggered
	require.True(watchFired(ws))

	// Ensure we don't get the objects back
	ws = memdb.NewWatchSet()
	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.Nil(out)

	ws = memdb.NewWatchSet()
	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
	require.NoError(err)
	require.Nil(out)

	// Ensure we see both policies
	iter, err := state.ScalingPoliciesByNamespace(ws, policy.Target[structs.ScalingTargetNamespace], "")
	require.NoError(err)
	count := 0
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
	}
	require.Equal(0, count)

	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.True(1001 == index)
	require.False(watchFired(ws))
}

func TestStateStore_StopJob_DeleteScalingPolicies(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	job := mock.Job()

	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	require.NoError(err)

	policy := mock.ScalingPolicy()
	policy.Target[structs.ScalingTargetJob] = job.ID
	err = state.UpsertScalingPolicies(1100, []*structs.ScalingPolicy{policy})
	require.NoError(err)

	// Ensure the scaling policy is present and start some watches
	wsGet := memdb.NewWatchSet()
	out, err := state.ScalingPolicyByTargetAndType(wsGet, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(out)
	wsList := memdb.NewWatchSet()
	_, err = state.ScalingPolicies(wsList)
	require.NoError(err)

	// Stop the job
	job, err = state.JobByID(nil, job.Namespace, job.ID)
	require.NoError(err)
	job.Stop = true
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1200, job)
	require.NoError(err)

	// Ensure:
	// * the scaling policy was deleted
	// * the watches were fired
	// * the table index was advanced
	require.True(watchFired(wsGet))
	require.True(watchFired(wsList))
	out, err = state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
	require.NoError(err)
	require.Nil(out)
	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.GreaterOrEqual(index, uint64(1200))
}

func TestStateStore_UnstopJob_UpsertScalingPolicies(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	job, policy := mock.JobWithScalingPolicy()
	job.Stop = true

	// establish watcher, verify there are no scaling policies yet
	ws := memdb.NewWatchSet()
	list, err := state.ScalingPolicies(ws)
	require.NoError(err)
	require.Nil(list.Next())

	// upsert a stopped job, verify that we don't fire the watcher or add any scaling policies
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	require.NoError(err)
	require.True(watchFired(ws))
	list, err = state.ScalingPolicies(ws)
	require.NoError(err)
	require.NotNil(list.Next())

	// Establish a new watchset
	ws = memdb.NewWatchSet()
	_, err = state.ScalingPolicies(ws)
	require.NoError(err)
	// Unstop this job, say you'll run it again...
	job.Stop = false
	err = state.UpsertJob(structs.MsgTypeTestSetup, 1100, job)
	require.NoError(err)

	// Ensure the scaling policy still exists, watch was not fired, index was not advanced
	out, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
	require.NoError(err)
	require.NotNil(out)
	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.EqualValues(index, 1000)
	require.False(watchFired(ws))
}

func TestStateStore_DeleteJob_DeleteScalingPolicies(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	job := mock.Job()

	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	require.NoError(err)

	policy := mock.ScalingPolicy()
	policy.Target[structs.ScalingTargetJob] = job.ID
	err = state.UpsertScalingPolicies(1001, []*structs.ScalingPolicy{policy})
	require.NoError(err)

	// Delete the job
	err = state.DeleteJob(1002, job.Namespace, job.ID)
	require.NoError(err)

	// Ensure the scaling policy was deleted
	ws := memdb.NewWatchSet()
	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
	require.NoError(err)
	require.Nil(out)
	index, err := state.Index("scaling_policy")
	require.NoError(err)
	require.True(index > 1001)
}

func TestStateStore_DeleteJob_DeleteScalingPoliciesPrefixBug(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	job := mock.Job()
	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1000, job))
	job2 := job.Copy()
	job2.ID = job.ID + "-but-longer"
	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2))

	policy := mock.ScalingPolicy()
	policy.Target[structs.ScalingTargetJob] = job.ID
	policy2 := mock.ScalingPolicy()
	policy2.Target[structs.ScalingTargetJob] = job2.ID
	require.NoError(state.UpsertScalingPolicies(1002, []*structs.ScalingPolicy{policy, policy2}))

	// Delete job with the shorter prefix-ID
	require.NoError(state.DeleteJob(1003, job.Namespace, job.ID))

	// Ensure only the associated scaling policy was deleted, not the one matching the job with the longer ID
	out, err := state.ScalingPolicyByID(nil, policy.ID)
	require.NoError(err)
	require.Nil(out)
	out, err = state.ScalingPolicyByID(nil, policy2.ID)
	require.NoError(err)
	require.NotNil(out)
}

// This test ensures that deleting a job that doesn't have any scaling policies
// will not cause the scaling_policy table index to increase, on either job
// registration or deletion.
func TestStateStore_DeleteJob_ScalingPolicyIndexNoop(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	job := mock.Job()

	prevIndex, err := state.Index("scaling_policy")
	require.NoError(err)

	err = state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
	require.NoError(err)

	newIndex, err := state.Index("scaling_policy")
	require.NoError(err)
	require.Equal(prevIndex, newIndex)

	// Delete the job
	err = state.DeleteJob(1002, job.Namespace, job.ID)
	require.NoError(err)

	newIndex, err = state.Index("scaling_policy")
	require.NoError(err)
	require.Equal(prevIndex, newIndex)
}

func TestStateStore_ScalingPoliciesByType(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	// Create scaling policies of different types
	pHorzA := mock.ScalingPolicy()
	pHorzA.Type = structs.ScalingPolicyTypeHorizontal
	pHorzB := mock.ScalingPolicy()
	pHorzB.Type = structs.ScalingPolicyTypeHorizontal

	pOther1 := mock.ScalingPolicy()
	pOther1.Type = "other-type-1"

	pOther2 := mock.ScalingPolicy()
	pOther2.Type = "other-type-2"

	// Create search routine
	search := func(t string) (found []string) {
		found = []string{}
		iter, err := state.ScalingPoliciesByTypePrefix(nil, t)
		require.NoError(err)

		for raw := iter.Next(); raw != nil; raw = iter.Next() {
			found = append(found, raw.(*structs.ScalingPolicy).Type)
		}
		return
	}

	// Create the policies
	var baseIndex uint64 = 1000
	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{pHorzA, pHorzB, pOther1, pOther2})
	require.NoError(err)

	// Check if we can read horizontal policies
	expect := []string{pHorzA.Type, pHorzB.Type}
	actual := search(structs.ScalingPolicyTypeHorizontal)
	require.ElementsMatch(expect, actual)

	// Check if we can read policies of other types
	expect = []string{pOther1.Type}
	actual = search("other-type-1")
	require.ElementsMatch(expect, actual)

	// Check that we can read policies by prefix
	expect = []string{"other-type-1", "other-type-2"}
	actual = search("other-type")
	require.Equal(expect, actual)

	// Check for empty result
	expect = []string{}
	actual = search("non-existing")
	require.ElementsMatch(expect, actual)
}

func TestStateStore_ScalingPoliciesByTypePrefix(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	// Create scaling policies of different types
	pHorzA := mock.ScalingPolicy()
	pHorzA.Type = structs.ScalingPolicyTypeHorizontal
	pHorzB := mock.ScalingPolicy()
	pHorzB.Type = structs.ScalingPolicyTypeHorizontal

	pOther1 := mock.ScalingPolicy()
	pOther1.Type = "other-type-1"

	pOther2 := mock.ScalingPolicy()
	pOther2.Type = "other-type-2"

	// Create search routine
	search := func(t string) (count int, found []string, err error) {
		found = []string{}
		iter, err := state.ScalingPoliciesByTypePrefix(nil, t)
		if err != nil {
			return
		}

		for raw := iter.Next(); raw != nil; raw = iter.Next() {
			count++
			found = append(found, raw.(*structs.ScalingPolicy).Type)
		}
		return
	}

	// Create the policies
	var baseIndex uint64 = 1000
	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{pHorzA, pHorzB, pOther1, pOther2})
	require.NoError(err)

	// Check if we can read horizontal policies
	expect := []string{pHorzA.Type, pHorzB.Type}
	count, found, err := search("h")

	sort.Strings(found)
	sort.Strings(expect)

	require.NoError(err)
	require.Equal(expect, found)
	require.Equal(2, count)

	// Check if we can read other prefix policies
	expect = []string{pOther1.Type, pOther2.Type}
	count, found, err = search("other")

	sort.Strings(found)
	sort.Strings(expect)

	require.NoError(err)
	require.Equal(expect, found)
	require.Equal(2, count)

	// Check for empty result
	expect = []string{}
	count, found, err = search("non-existing")

	sort.Strings(found)
	sort.Strings(expect)

	require.NoError(err)
	require.Equal(expect, found)
	require.Equal(0, count)
}

func TestStateStore_ScalingPoliciesByJob(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)
	policyA := mock.ScalingPolicy()
	policyB1 := mock.ScalingPolicy()
	policyB2 := mock.ScalingPolicy()
	policyB1.Target[structs.ScalingTargetJob] = policyB2.Target[structs.ScalingTargetJob]

	// Create the policies
	var baseIndex uint64 = 1000
	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policyA, policyB1, policyB2})
	require.NoError(err)

	iter, err := state.ScalingPoliciesByJob(nil,
		policyA.Target[structs.ScalingTargetNamespace],
		policyA.Target[structs.ScalingTargetJob], "")
	require.NoError(err)

	// Ensure we see expected policies
	count := 0
	found := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
		found = append(found, raw.(*structs.ScalingPolicy).Target[structs.ScalingTargetGroup])
	}
	require.Equal(1, count)
	sort.Strings(found)
	expect := []string{policyA.Target[structs.ScalingTargetGroup]}
	sort.Strings(expect)
	require.Equal(expect, found)

	iter, err = state.ScalingPoliciesByJob(nil,
		policyB1.Target[structs.ScalingTargetNamespace],
		policyB1.Target[structs.ScalingTargetJob], "")
	require.NoError(err)

	// Ensure we see expected policies
	count = 0
	found = []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
		found = append(found, raw.(*structs.ScalingPolicy).Target[structs.ScalingTargetGroup])
	}
	require.Equal(2, count)
	sort.Strings(found)
	expect = []string{
		policyB1.Target[structs.ScalingTargetGroup],
		policyB2.Target[structs.ScalingTargetGroup],
	}
	sort.Strings(expect)
	require.Equal(expect, found)
}

func TestStateStore_ScalingPoliciesByJob_PrefixBug(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	jobPrefix := "job-name-" + uuid.Generate()

	state := testStateStore(t)
	policy1 := mock.ScalingPolicy()
	policy1.Target[structs.ScalingTargetJob] = jobPrefix
	policy2 := mock.ScalingPolicy()
	policy2.Target[structs.ScalingTargetJob] = jobPrefix + "-more"

	// Create the policies
	var baseIndex uint64 = 1000
	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policy1, policy2})
	require.NoError(err)

	iter, err := state.ScalingPoliciesByJob(nil,
		policy1.Target[structs.ScalingTargetNamespace],
		jobPrefix, "")
	require.NoError(err)

	// Ensure we see expected policies
	count := 0
	found := []string{}
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		count++
		found = append(found, raw.(*structs.ScalingPolicy).ID)
	}
	require.Equal(1, count)
	expect := []string{policy1.ID}
	require.Equal(expect, found)
}

func TestStateStore_ScalingPolicyByTargetAndType(t *testing.T) {
	t.Parallel()

	require := require.New(t)

	state := testStateStore(t)

	// Create scaling policies
	policyA := mock.ScalingPolicy()
	// Same target, different type
	policyB := mock.ScalingPolicy()
	policyC := mock.ScalingPolicy()
	for k, v := range policyB.Target {
		policyC.Target[k] = v
	}
	policyC.Type = "other-type"

	// Create the policies
	var baseIndex uint64 = 1000
	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policyA, policyB, policyC})
	require.NoError(err)

	// Check if we can retrieve the right policies
	found, err := state.ScalingPolicyByTargetAndType(nil, policyA.Target, policyA.Type)
	require.NoError(err)
	require.Equal(policyA, found)

	// Check for wrong type
	found, err = state.ScalingPolicyByTargetAndType(nil, policyA.Target, "wrong_type")
	require.NoError(err)
	require.Nil(found)

	// Check for same target but different type
	found, err = state.ScalingPolicyByTargetAndType(nil, policyB.Target, policyB.Type)
	require.NoError(err)
	require.Equal(policyB, found)

	found, err = state.ScalingPolicyByTargetAndType(nil, policyB.Target, policyC.Type)
	require.NoError(err)
	require.Equal(policyC, found)
}

func TestStateStore_UpsertScalingEvent(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	job := mock.Job()
	groupName := job.TaskGroups[0].Name

	newEvent := structs.NewScalingEvent("message 1").SetMeta(map[string]interface{}{
		"a": 1,
	})

	wsAll := memdb.NewWatchSet()
	all, err := state.ScalingEvents(wsAll)
	require.NoError(err)
	require.Nil(all.Next())

	ws := memdb.NewWatchSet()
	out, _, err := state.ScalingEventsByJob(ws, job.Namespace, job.ID)
	require.NoError(err)
	require.Nil(out)

	err = state.UpsertScalingEvent(1000, &structs.ScalingEventRequest{
		Namespace:    job.Namespace,
		JobID:        job.ID,
		TaskGroup:    groupName,
		ScalingEvent: newEvent,
	})
	require.NoError(err)
	require.True(watchFired(ws))
	require.True(watchFired(wsAll))

	ws = memdb.NewWatchSet()
	out, eventsIndex, err := state.ScalingEventsByJob(ws, job.Namespace, job.ID)
	require.NoError(err)
	require.Equal(map[string][]*structs.ScalingEvent{
		groupName: {newEvent},
	}, out)
	require.EqualValues(eventsIndex, 1000)

	iter, err := state.ScalingEvents(ws)
	require.NoError(err)

	count := 0
	jobsReturned := []string{}
	var jobEvents *structs.JobScalingEvents
	for {
		raw := iter.Next()
		if raw == nil {
			break
		}
		jobEvents = raw.(*structs.JobScalingEvents)
		jobsReturned = append(jobsReturned, jobEvents.JobID)
		count++
	}
	require.Equal(1, count)
	require.EqualValues(jobEvents.ModifyIndex, 1000)
	require.EqualValues(jobEvents.ScalingEvents[groupName][0].CreateIndex, 1000)

	index, err := state.Index("scaling_event")
	require.NoError(err)
	require.ElementsMatch([]string{job.ID}, jobsReturned)
	require.Equal(map[string][]*structs.ScalingEvent{
		groupName: {newEvent},
	}, jobEvents.ScalingEvents)
	require.EqualValues(1000, index)
	require.False(watchFired(ws))
}

func TestStateStore_UpsertScalingEvent_LimitAndOrder(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	namespace := uuid.Generate()
	jobID := uuid.Generate()
	group1 := uuid.Generate()
	group2 := uuid.Generate()

	index := uint64(1000)
	for i := 1; i <= structs.JobTrackedScalingEvents+10; i++ {
		newEvent := structs.NewScalingEvent("").SetMeta(map[string]interface{}{
			"i":     i,
			"group": group1,
		})
		err := state.UpsertScalingEvent(index, &structs.ScalingEventRequest{
			Namespace:    namespace,
			JobID:        jobID,
			TaskGroup:    group1,
			ScalingEvent: newEvent,
		})
		index++
		require.NoError(err)

		newEvent = structs.NewScalingEvent("").SetMeta(map[string]interface{}{
			"i":     i,
			"group": group2,
		})
		err = state.UpsertScalingEvent(index, &structs.ScalingEventRequest{
			Namespace:    namespace,
			JobID:        jobID,
			TaskGroup:    group2,
			ScalingEvent: newEvent,
		})
		index++
		require.NoError(err)
	}

	out, _, err := state.ScalingEventsByJob(nil, namespace, jobID)
	require.NoError(err)
	require.Len(out, 2)

	expectedEvents := []int{}
	for i := structs.JobTrackedScalingEvents; i > 0; i-- {
		expectedEvents = append(expectedEvents, i+10)
	}

	// checking order and content
	require.Len(out[group1], structs.JobTrackedScalingEvents)
	actualEvents := []int{}
	for _, event := range out[group1] {
		require.Equal(group1, event.Meta["group"])
		actualEvents = append(actualEvents, event.Meta["i"].(int))
	}
	require.Equal(expectedEvents, actualEvents)

	// checking order and content
	require.Len(out[group2], structs.JobTrackedScalingEvents)
	actualEvents = []int{}
	for _, event := range out[group2] {
		require.Equal(group2, event.Meta["group"])
		actualEvents = append(actualEvents, event.Meta["i"].(int))
	}
	require.Equal(expectedEvents, actualEvents)
}

func TestStateStore_RestoreScalingEvents(t *testing.T) {
	t.Parallel()
	require := require.New(t)

	state := testStateStore(t)
	jobScalingEvents := &structs.JobScalingEvents{
		Namespace: uuid.Generate(),
		JobID:     uuid.Generate(),
		ScalingEvents: map[string][]*structs.ScalingEvent{
			uuid.Generate(): {
				structs.NewScalingEvent(uuid.Generate()),
			},
		},
	}

	restore, err := state.Restore()
	require.NoError(err)

	err = restore.ScalingEventsRestore(jobScalingEvents)
	require.NoError(err)
	require.NoError(restore.Commit())

	ws := memdb.NewWatchSet()
	out, _, err := state.ScalingEventsByJob(ws, jobScalingEvents.Namespace,
		jobScalingEvents.JobID)
	require.NoError(err)
	require.NotNil(out)
	require.EqualValues(jobScalingEvents.ScalingEvents, out)
}

func TestStateStore_Abandon(t *testing.T) {
	t.Parallel()

	s := testStateStore(t)
	abandonCh := s.AbandonCh()
	s.Abandon()
	select {
	case <-abandonCh:
	default:
		t.Fatalf("bad")
	}
}

// Verifies that an error is returned when an allocation doesn't exist in the state store.
func TestStateSnapshot_DenormalizeAllocationDiffSlice_AllocDoesNotExist(t *testing.T) {
	t.Parallel()

	state := testStateStore(t)
	alloc := mock.Alloc()
	require := require.New(t)

	// Insert job
	err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job)
	require.NoError(err)

	allocDiffs := []*structs.AllocationDiff{
		{
			ID: alloc.ID,
		},
	}

	snap, err := state.Snapshot()
	require.NoError(err)

	denormalizedAllocs, err := snap.DenormalizeAllocationDiffSlice(allocDiffs)

	require.EqualError(err, fmt.Sprintf("alloc %v doesn't exist", alloc.ID))
	require.Nil(denormalizedAllocs)
}

// TestStateStore_SnapshotMinIndex_OK asserts StateStore.SnapshotMinIndex blocks
// until the StateStore's latest index is >= the requested index.
func TestStateStore_SnapshotMinIndex_OK(t *testing.T) {
	t.Parallel()

	s := testStateStore(t)
	index, err := s.LatestIndex()
	require.NoError(t, err)

	node := mock.Node()
	require.NoError(t, s.UpsertNode(structs.MsgTypeTestSetup, index+1, node))

	// Assert SnapshotMinIndex returns immediately if index < latest index
	ctx, cancel := context.WithTimeout(context.Background(), 0)
	snap, err := s.SnapshotMinIndex(ctx, index)
	cancel()
	require.NoError(t, err)

	snapIndex, err := snap.LatestIndex()
	require.NoError(t, err)
	if snapIndex <= index {
		require.Fail(t, "snapshot index should be greater than index")
	}

	// Assert SnapshotMinIndex returns immediately if index == latest index
	ctx, cancel = context.WithTimeout(context.Background(), 0)
	snap, err = s.SnapshotMinIndex(ctx, index+1)
	cancel()
	require.NoError(t, err)

	snapIndex, err = snap.LatestIndex()
	require.NoError(t, err)
	require.Equal(t, snapIndex, index+1)

	// Assert SnapshotMinIndex blocks if index > latest index
	errCh := make(chan error, 1)
	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	go func() {
		defer close(errCh)
		waitIndex := index + 2
		snap, err := s.SnapshotMinIndex(ctx, waitIndex)
		if err != nil {
			errCh <- err
			return
		}

		snapIndex, err := snap.LatestIndex()
		if err != nil {
			errCh <- err
			return
		}

		if snapIndex < waitIndex {
			errCh <- fmt.Errorf("snapshot index < wait index: %d < %d", snapIndex, waitIndex)
			return
		}
	}()

	select {
	case err := <-errCh:
		require.NoError(t, err)
	case <-time.After(500 * time.Millisecond):
		// Let it block for a bit before unblocking by upserting
	}

	node.Name = "hal"
	require.NoError(t, s.UpsertNode(structs.MsgTypeTestSetup, index+2, node))

	select {
	case err := <-errCh:
		require.NoError(t, err)
	case <-time.After(5 * time.Second):
		require.Fail(t, "timed out waiting for SnapshotMinIndex to unblock")
	}
}

// TestStateStore_SnapshotMinIndex_Timeout asserts StateStore.SnapshotMinIndex
// returns an error if the desired index is not reached within the deadline.
func TestStateStore_SnapshotMinIndex_Timeout(t *testing.T) {
	t.Parallel()

	s := testStateStore(t)
	index, err := s.LatestIndex()
	require.NoError(t, err)

	// Assert SnapshotMinIndex blocks if index > latest index
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	snap, err := s.SnapshotMinIndex(ctx, index+1)
	require.EqualError(t, err, context.DeadlineExceeded.Error())
	require.Nil(t, snap)
}

// watchFired is a helper for unit tests that returns if the given watch set
// fired (it doesn't care which watch actually fired). This uses a fixed
// timeout since we already expect the event happened before calling this and
// just need to distinguish a fire from a timeout. We do need a little time to
// allow the watch to set up any goroutines, though.
func watchFired(ws memdb.WatchSet) bool {
	timedOut := ws.Watch(time.After(50 * time.Millisecond))
	return !timedOut
}

// NodeIDSort is used to sort nodes by ID
type NodeIDSort []*structs.Node

func (n NodeIDSort) Len() int {
	return len(n)
}

func (n NodeIDSort) Less(i, j int) bool {
	return n[i].ID < n[j].ID
}

func (n NodeIDSort) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}

// JobIDis used to sort jobs by id
type JobIDSort []*structs.Job

func (n JobIDSort) Len() int {
	return len(n)
}

func (n JobIDSort) Less(i, j int) bool {
	return n[i].ID < n[j].ID
}

func (n JobIDSort) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}

// EvalIDis used to sort evals by id
type EvalIDSort []*structs.Evaluation

func (n EvalIDSort) Len() int {
	return len(n)
}

func (n EvalIDSort) Less(i, j int) bool {
	return n[i].ID < n[j].ID
}

func (n EvalIDSort) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}

// AllocIDsort used to sort allocations by id
type AllocIDSort []*structs.Allocation

func (n AllocIDSort) Len() int {
	return len(n)
}

func (n AllocIDSort) Less(i, j int) bool {
	return n[i].ID < n[j].ID
}

func (n AllocIDSort) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}