add state store test to ensure PlacedCanaries is updated

nomad/state/state_store_test.go

@@ -6382,6 +6382,109 @@ func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_Nonexistent(t *testing.
 	}
 }
 
+// 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(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(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(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(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(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(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) {

scheduler/generic_sched_test.go

@@ -2029,17 +2029,28 @@ func TestServiceSched_JobModify_Canaries(t *testing.T) {
 	if plan.Deployment == nil {
 		t.Fatalf("bad: %#v", plan)
 	}
-	state, ok := plan.Deployment.TaskGroups[job.TaskGroups[0].Name]
-	if !ok {
-		t.Fatalf("bad: %#v", plan)
-	}
-	if state.DesiredTotal != 10 && state.DesiredCanaries != desiredUpdates {
-		t.Fatalf("bad: %#v", state)
-	}
+
+	// Ensure local state was not altered in scheduler
+	staleState, ok := plan.Deployment.TaskGroups[job.TaskGroups[0].Name]
+	require.True(t, ok)
+
+	require.Equal(t, 0, len(staleState.PlacedCanaries))
+
+	ws := memdb.NewWatchSet()
+
+	// Grab the latest state
+	deploy, err := h.State.DeploymentByID(ws, plan.Deployment.ID)
+	require.NoError(t, err)
+
+	state, ok := deploy.TaskGroups[job.TaskGroups[0].Name]
+	require.True(t, ok)
+
+	require.Equal(t, 10, state.DesiredTotal)
+	require.Equal(t, state.DesiredCanaries, desiredUpdates)
 
 	// Assert the canaries were added to the placed list
 	if len(state.PlacedCanaries) != desiredUpdates {
-		t.Fatalf("bad: %#v", state)
+		assert.Fail(t, "expected PlacedCanaries to equal desiredUpdates", state)
 	}
 }