Fix flaky transform_test.go

Since snapshots include query references and query plan references,
and since those can be orderered nondeterministically, the test can
fail if the snapshot is built with fields in an unexpected order.

The existing test tries to account for that by enumerating all
possible combinations, but this is very difficult to maintain, and
very difficult to scale to additional fields with this pattern.

Instead, rewrite snapshots to a canonical format to ensure comparisons
are stable. We can't just sort the array fields before serializing the
snapshot for comparison, since QueryReferences and similar items are
referenced by an index into their array. To work around that, we sort
the QueryReferences (tracking their original index), and rewrite all
entries that reference them to reference their new index instead.

See discussion in #630 [1].

[1]: #630 (comment)

output/transform/transform_test.go

@@ -3,6 +3,8 @@ package transform_test
 import (
 	"encoding/binary"
 	"encoding/json"
+	"slices"
+	"sort"
 	"testing"
 	"time"
 
@@ -46,9 +48,9 @@ func TestStatements(t *testing.T) {
 	diffState.PlanStats[pKey2] = state.DiffedPostgresStatementStats{Calls: 24}
 
 	actual := transform.StateToSnapshot(newState, diffState, transientState)
+	makeCanonical(actual)
 	actualJSON, _ := json.Marshal(actual)
 
-	// Query: 0, 1, Plan: 0, 1
 	expected := pganalyze_collector.FullSnapshot{
 		Config:             &pganalyze_collector.CollectorConfig{},
 		CollectorStatistic: &pganalyze_collector.CollectorStatistic{},
@@ -66,189 +68,198 @@ func TestStatements(t *testing.T) {
 		ServerStatistic: &pganalyze_collector.ServerStatistic{},
 		Replication:     &pganalyze_collector.Replication{},
 		QueryReferences: []*pganalyze_collector.QueryReference{
-			&pganalyze_collector.QueryReference{
+			{
 				DatabaseIdx: 0,
 				RoleIdx:     0,
 				Fingerprint: fpBuf1,
 			},
-			&pganalyze_collector.QueryReference{
+			{
 				DatabaseIdx: 0,
 				RoleIdx:     0,
 				Fingerprint: fpBuf2,
 			},
 		},
 		QueryInformations: []*pganalyze_collector.QueryInformation{
-			&pganalyze_collector.QueryInformation{
+			{
 				QueryIdx:        0,
 				NormalizedQuery: "SELECT 1",
 				QueryIds:        []int64{1},
 			},
-			&pganalyze_collector.QueryInformation{
+			{
 				QueryIdx:        1,
 				NormalizedQuery: "SELECT * FROM test",
 				QueryIds:        []int64{2},
 			},
 		},
 		QueryStatistics: []*pganalyze_collector.QueryStatistic{
-			&pganalyze_collector.QueryStatistic{
+			{
 				QueryIdx: 0,
 				Calls:    1,
 			},
-			&pganalyze_collector.QueryStatistic{
+			{
 				QueryIdx: 1,
 				Calls:    13,
 			},
 		},
 		QueryPlanReferences: []*pganalyze_collector.QueryPlanReference{
-			&pganalyze_collector.QueryPlanReference{
+			{
 				QueryIdx:       1,
 				OriginalPlanId: 111,
 			},
-			&pganalyze_collector.QueryPlanReference{
+			{
 				QueryIdx:       1,
 				OriginalPlanId: 222,
 			},
 		},
 		QueryPlanInformations: []*pganalyze_collector.QueryPlanInformation{
-			&pganalyze_collector.QueryPlanInformation{
+			{
 				QueryPlanIdx:     0,
 				ExplainPlan:      "Index Scan",
 				PlanCapturedTime: timestamppb.New(capturedTime),
 			},
-			&pganalyze_collector.QueryPlanInformation{
+			{
 				QueryPlanIdx:     1,
 				ExplainPlan:      "Bitmap Heap Scan",
 				PlanCapturedTime: timestamppb.New(capturedTime),
 			},
 		},
 		QueryPlanStatistics: []*pganalyze_collector.QueryPlanStatistic{
-			&pganalyze_collector.QueryPlanStatistic{
+			{
 				QueryPlanIdx: 0,
 				Calls:        2,
 			},
-			&pganalyze_collector.QueryPlanStatistic{
+			{
 				QueryPlanIdx: 1,
 				Calls:        24,
 			},
 		},
 	}
+	makeCanonical(expected)
 	expectedJSON, _ := json.Marshal(expected)
 
-	// Query: 1, 0, Plan: 0, 1
-	var expectedAlt pganalyze_collector.FullSnapshot
-	json.Unmarshal(expectedJSON, &expectedAlt)
-	expectedAlt.QueryReferences = []*pganalyze_collector.QueryReference{
-		&pganalyze_collector.QueryReference{
-			DatabaseIdx: 0,
-			RoleIdx:     0,
-			Fingerprint: fpBuf2,
-		},
-		&pganalyze_collector.QueryReference{
-			DatabaseIdx: 0,
-			RoleIdx:     0,
-			Fingerprint: fpBuf1,
-		},
-	}
-	expectedAlt.QueryInformations = []*pganalyze_collector.QueryInformation{
-		&pganalyze_collector.QueryInformation{
-			QueryIdx:        0,
-			NormalizedQuery: "SELECT * FROM test",
-			QueryIds:        []int64{2},
-		},
-		&pganalyze_collector.QueryInformation{
-			QueryIdx:        1,
-			NormalizedQuery: "SELECT 1",
-			QueryIds:        []int64{1},
-		},
-	}
-	expectedAlt.QueryStatistics = []*pganalyze_collector.QueryStatistic{
-		&pganalyze_collector.QueryStatistic{
-			QueryIdx: 0,
-			Calls:    13,
-		},
-		&pganalyze_collector.QueryStatistic{
-			QueryIdx: 1,
-			Calls:    1,
-		},
+	if string(expectedJSON) != string(actualJSON) {
+		t.Errorf("\nExpected:%+v\n\tActual: %+v\n\n", string(expectedJSON), string(actualJSON))
 	}
-	expectedJSONAlt, _ := json.Marshal(expectedAlt)
+}
 
-	// Query: 1, 0, Plan: 1, 0
-	expectedAlt.QueryPlanReferences = []*pganalyze_collector.QueryPlanReference{
-		&pganalyze_collector.QueryPlanReference{
-			QueryIdx:       1,
-			OriginalPlanId: 222,
-		},
-		&pganalyze_collector.QueryPlanReference{
-			QueryIdx:       1,
-			OriginalPlanId: 111,
-		},
-	}
-	expectedAlt.QueryPlanInformations = []*pganalyze_collector.QueryPlanInformation{
-		&pganalyze_collector.QueryPlanInformation{
-			QueryPlanIdx:     0,
-			ExplainPlan:      "Bitmap Heap Scan",
-			PlanCapturedTime: timestamppb.New(capturedTime),
-		},
-		&pganalyze_collector.QueryPlanInformation{
-			QueryPlanIdx:     1,
-			ExplainPlan:      "Index Scan",
-			PlanCapturedTime: timestamppb.New(capturedTime),
-		},
-	}
-	expectedAlt.QueryPlanStatistics = []*pganalyze_collector.QueryPlanStatistic{
-		&pganalyze_collector.QueryPlanStatistic{
-			QueryPlanIdx: 0,
-			Calls:        24,
-		},
-		&pganalyze_collector.QueryPlanStatistic{
-			QueryPlanIdx: 1,
-			Calls:        2,
-		},
-	}
-	expectedJSONAlt2, _ := json.Marshal(expectedAlt)
-
-	// Query: 0, 1, Plan: 1, 0
-	expectedAlt.QueryReferences = []*pganalyze_collector.QueryReference{
-		&pganalyze_collector.QueryReference{
-			DatabaseIdx: 0,
-			RoleIdx:     0,
-			Fingerprint: fpBuf1,
-		},
-		&pganalyze_collector.QueryReference{
-			DatabaseIdx: 0,
-			RoleIdx:     0,
-			Fingerprint: fpBuf2,
-		},
+type OriginalQueryRef struct {
+	Original    *pganalyze_collector.QueryReference
+	OriginalIdx int32
+}
+
+type OriginalPlanRef struct {
+	Original    *pganalyze_collector.QueryPlanReference
+	OriginalIdx int32
+}
+
+// Takes a snapshot and rewrites it to a canonical form, with QueryReferences
+// and similar fields sorted into a single canonical order.
+func makeCanonical(snapshot pganalyze_collector.FullSnapshot) {
+	// ensure query references occur in a consistent order
+	queryRefs := make([]*OriginalQueryRef, len(snapshot.QueryReferences))
+	for i, qRef := range snapshot.QueryReferences {
+		queryRefs[i] = &OriginalQueryRef{
+			Original:    qRef,
+			OriginalIdx: int32(i),
+		}
 	}
-	expectedAlt.QueryInformations = []*pganalyze_collector.QueryInformation{
-		&pganalyze_collector.QueryInformation{
-			QueryIdx:        0,
-			NormalizedQuery: "SELECT 1",
-			QueryIds:        []int64{1},
-		},
-		&pganalyze_collector.QueryInformation{
-			QueryIdx:        1,
-			NormalizedQuery: "SELECT * FROM test",
-			QueryIds:        []int64{2},
-		},
+
+	sort.Slice(queryRefs, func(i, j int) bool {
+		a := queryRefs[i].Original
+		b := queryRefs[j].Original
+
+		if a.DatabaseIdx < b.DatabaseIdx {
+			return true
+		}
+		if a.DatabaseIdx > b.DatabaseIdx {
+			return false
+		}
+		if a.RoleIdx < b.RoleIdx {
+			return true
+		}
+		if a.RoleIdx > b.RoleIdx {
+			return true
+		}
+		if len(a.Fingerprint) < len(b.Fingerprint) {
+			return true
+		}
+		if len(a.Fingerprint) > len(b.Fingerprint) {
+			return false
+		}
+		for i, aFpByte := range a.Fingerprint {
+			bFpByte := b.Fingerprint[i]
+			if aFpByte < bFpByte {
+				return true
+			}
+			if aFpByte > bFpByte {
+				return false
+			}
+		}
+		return false
+	})
+
+	for i, qRef := range queryRefs {
+		snapshot.QueryReferences[i] = qRef.Original
+
+		qInfoIdx := slices.IndexFunc(snapshot.QueryInformations, func(item *pganalyze_collector.QueryInformation) bool {
+			return item.QueryIdx == qRef.OriginalIdx
+		})
+		qInfo := snapshot.QueryInformations[qInfoIdx]
+		qInfo.QueryIdx = int32(i)
+
+		qStatsIdx := slices.IndexFunc(snapshot.QueryStatistics, func(item *pganalyze_collector.QueryStatistic) bool {
+			return item.QueryIdx == qRef.OriginalIdx
+		})
+		qStats := snapshot.QueryStatistics[qStatsIdx]
+		qStats.QueryIdx = int32(i)
 	}
-	expectedAlt.QueryStatistics = []*pganalyze_collector.QueryStatistic{
-		&pganalyze_collector.QueryStatistic{
-			QueryIdx: 0,
-			Calls:    1,
-		},
-		&pganalyze_collector.QueryStatistic{
-			QueryIdx: 1,
-			Calls:    13,
-		},
+	sort.Slice(snapshot.QueryInformations, func(i, j int) bool {
+		a := snapshot.QueryInformations[i]
+		b := snapshot.QueryInformations[j]
+		return a.QueryIdx < b.QueryIdx
+	})
+	sort.Slice(snapshot.QueryStatistics, func(i, j int) bool {
+		a := snapshot.QueryStatistics[i]
+		b := snapshot.QueryStatistics[j]
+		return a.QueryIdx < b.QueryIdx
+	})
+
+	// ensure plan references occur in a consistent order
+	planRefs := make([]OriginalPlanRef, len(snapshot.QueryPlanReferences))
+	for i, planRef := range snapshot.QueryPlanReferences {
+		planRefs[i] = OriginalPlanRef{
+			Original:    planRef,
+			OriginalIdx: int32(i),
+		}
 	}
-	expectedJSONAlt3, _ := json.Marshal(expectedAlt)
 
-	if string(expectedJSON) != string(actualJSON) &&
-		string(expectedJSONAlt) != string(actualJSON) &&
-		string(expectedJSONAlt2) != string(actualJSON) &&
-		string(expectedJSONAlt3) != string(actualJSON) {
-		t.Errorf("\nExpected:%+v\n\tActual: %+v\n\n", string(expectedJSON), string(actualJSON))
+	sort.Slice(planRefs, func(i, j int) bool {
+		return planRefs[i].Original.OriginalPlanId < planRefs[j].Original.OriginalPlanId
+	})
+
+	for i, planRef := range planRefs {
+		snapshot.QueryPlanReferences[i] = planRef.Original
+
+		pInfoIdx := slices.IndexFunc(snapshot.QueryPlanInformations, func(item *pganalyze_collector.QueryPlanInformation) bool {
+			return item.QueryPlanIdx == planRef.OriginalIdx
+		})
+		pInfo := snapshot.QueryPlanInformations[pInfoIdx]
+		pInfo.QueryPlanIdx = int32(i)
+
+		pStatsIdx := slices.IndexFunc(snapshot.QueryPlanStatistics, func(item *pganalyze_collector.QueryPlanStatistic) bool {
+			return item.QueryPlanIdx == planRef.OriginalIdx
+		})
+		pStats := snapshot.QueryPlanStatistics[pStatsIdx]
+		pStats.QueryPlanIdx = int32(i)
 	}
+	sort.Slice(snapshot.QueryPlanInformations, func(i, j int) bool {
+		a := snapshot.QueryPlanInformations[i]
+		b := snapshot.QueryPlanInformations[j]
+		return a.QueryPlanIdx < b.QueryPlanIdx
+	})
+	sort.Slice(snapshot.QueryPlanStatistics, func(i, j int) bool {
+		a := snapshot.QueryPlanStatistics[i]
+		b := snapshot.QueryPlanStatistics[j]
+		return a.QueryPlanIdx < b.QueryPlanIdx
+	})
 }