praefect: Explicitly track expected generation numbers

The coordinator tests which verify replication logic for transactional
mutators currently computes the expected generation number ad-hoc. The
logic is about to change though and become a bit more complex. As a
preparatory change, this commit thus explicitly records the expected
generation number to make it easier to change.

1 files changed, 27 insertions, 30 deletions

diff --git a/internal/praefect/coordinator_pg_test.go b/internal/praefect/coordinator_pg_test.go
index 574c433b3..419ad8962 100644
--- a/internal/praefect/coordinator_pg_test.go
+++ b/internal/praefect/coordinator_pg_test.go
@@ -32,12 +32,13 @@ func getDB(t *testing.T) glsql.DB {
 
 func TestStreamDirectorMutator_Transaction(t *testing.T) {
 	type node struct {
-		primary           bool
-		vote              string
-		shouldSucceed     bool
-		shouldGetRepl     bool
-		shouldParticipate bool
-		generation        int
+		primary            bool
+		vote               string
+		shouldSucceed      bool
+		shouldGetRepl      bool
+		shouldParticipate  bool
+		generation         int
+		expectedGeneration int
 	}
 
 	testcases := []struct {
@@ -47,48 +48,48 @@ func TestStreamDirectorMutator_Transaction(t *testing.T) {
 		{
 			desc: "successful vote should not create replication jobs",
 			nodes: []node{
-				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
+				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
+				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
+				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
 			},
 		},
 		{
 			desc: "failing vote should not create replication jobs",
 			nodes: []node{
-				{primary: true, vote: "foo", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "qux", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "bar", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true},
+				{primary: true, vote: "foo", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 0},
+				{primary: false, vote: "qux", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 0},
+				{primary: false, vote: "bar", shouldSucceed: false, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 0},
 			},
 		},
 		{
 			desc: "primary should reach quorum with disagreeing secondary",
 			nodes: []node{
-				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "barfoo", shouldSucceed: false, shouldGetRepl: true, shouldParticipate: true},
+				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
+				{primary: false, vote: "barfoo", shouldSucceed: false, shouldGetRepl: true, shouldParticipate: true, expectedGeneration: 0},
 			},
 		},
 		{
 			desc: "quorum should create replication jobs for disagreeing node",
 			nodes: []node{
-				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true},
-				{primary: false, vote: "barfoo", shouldSucceed: false, shouldGetRepl: true, shouldParticipate: true},
+				{primary: true, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
+				{primary: false, vote: "foobar", shouldSucceed: true, shouldGetRepl: false, shouldParticipate: true, expectedGeneration: 1},
+				{primary: false, vote: "barfoo", shouldSucceed: false, shouldGetRepl: true, shouldParticipate: true, expectedGeneration: 0},
 			},
 		},
 		{
 			desc: "only consistent secondaries should participate",
 			nodes: []node{
-				{primary: true, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: 1},
-				{primary: false, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: 1},
-				{shouldParticipate: false, shouldGetRepl: true, generation: 0},
-				{shouldParticipate: false, shouldGetRepl: true, generation: datastore.GenerationUnknown},
+				{primary: true, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: 1, expectedGeneration: 2},
+				{primary: false, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: 1, expectedGeneration: 2},
+				{shouldParticipate: false, shouldGetRepl: true, generation: 0, expectedGeneration: 0},
+				{shouldParticipate: false, shouldGetRepl: true, generation: datastore.GenerationUnknown, expectedGeneration: datastore.GenerationUnknown},
 			},
 		},
 		{
 			desc: "secondaries should not participate when primary's generation is unknown",
 			nodes: []node{
-				{primary: true, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: datastore.GenerationUnknown},
-				{shouldParticipate: false, shouldGetRepl: true, generation: datastore.GenerationUnknown},
+				{primary: true, vote: "foobar", shouldSucceed: true, shouldParticipate: true, generation: datastore.GenerationUnknown, expectedGeneration: 0},
+				{shouldParticipate: false, shouldGetRepl: true, generation: datastore.GenerationUnknown, expectedGeneration: datastore.GenerationUnknown},
 			},
 		},
 		{
@@ -97,8 +98,8 @@ func TestStreamDirectorMutator_Transaction(t *testing.T) {
 			// replication jobs.
 			desc: "unstarted transaction should create replication jobs",
 			nodes: []node{
-				{primary: true, shouldSucceed: true, shouldGetRepl: false},
-				{primary: false, shouldSucceed: false, shouldGetRepl: true},
+				{primary: true, shouldSucceed: true, shouldGetRepl: false, expectedGeneration: 1},
+				{primary: false, shouldSucceed: false, shouldGetRepl: true, expectedGeneration: 0},
 			},
 		},
 	}
@@ -235,11 +236,7 @@ func TestStreamDirectorMutator_Transaction(t *testing.T) {
 			for i, n := range tc.nodes {
 				gen, err := rs.GetGeneration(ctx, repo.StorageName, repo.RelativePath, storageNodes[i].Storage)
 				require.NoError(t, err)
-				expectedGeneration := n.generation
-				if n.shouldSucceed {
-					expectedGeneration++
-				}
-				require.Equal(t, expectedGeneration, gen)
+				require.Equal(t, n.expectedGeneration, gen, "node %d has wrong generation", i)
 			}
 
 			replicationWaitGroup.Wait()