Maybe fix send-to-device tests

syncapi/storage/storage_test.go

@@ -60,6 +60,17 @@ func TestWriteEvents(t *testing.T) {
 	})
 }
 
+func WithSnapshot(t *testing.T, db storage.Database, f func(snapshot storage.DatabaseSnapshot)) {
+	snapshot, err := db.NewDatabaseSnapshot(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	f(snapshot)
+	if err := snapshot.Rollback(); err != nil {
+		t.Fatal(err)
+	}
+}
+
 // These tests assert basic functionality of RecentEvents for PDUs
 func TestRecentEventsPDU(t *testing.T) {
 	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
@@ -79,16 +90,13 @@ func TestRecentEventsPDU(t *testing.T) {
 		// dummy room to make sure SQL queries are filtering on room ID
 		MustWriteEvents(t, db, test.NewRoom(t, alice).Events())
 
-		snapshot, err := db.NewDatabaseSnapshot(ctx)
+		var latest types.StreamPosition
-		if err != nil {
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
-			t.Fatal(err)
+			var err error
-		}
+			if latest, err = snapshot.MaxStreamPositionForPDUs(ctx); err != nil {
-		defer snapshot.Rollback() // nolint:errcheck
+				t.Fatal("failed to get MaxStreamPositionForPDUs: %w", err)
-
-		latest, err := snapshot.MaxStreamPositionForPDUs(ctx)
-		if err != nil {
-			t.Fatalf("failed to get MaxStreamPositionForPDUs: %s", err)
 			}
+		})
 
 		testCases := []struct {
 			Name         string
@@ -146,14 +154,19 @@ func TestRecentEventsPDU(t *testing.T) {
 			tc := testCases[i]
 			t.Run(tc.Name, func(st *testing.T) {
 				var filter gomatrixserverlib.RoomEventFilter
+				var gotEvents []types.StreamEvent
+				var limited bool
 				filter.Limit = tc.Limit
-				gotEvents, limited, err := snapshot.RecentEvents(ctx, r.ID, types.Range{
+				WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
+					var err error
+					gotEvents, limited, err = snapshot.RecentEvents(ctx, r.ID, types.Range{
 						From: tc.From,
 						To:   tc.To,
 					}, &filter, !tc.ReverseOrder, true)
 					if err != nil {
 						st.Fatalf("failed to do sync: %s", err)
 					}
+				})
 				if limited != tc.WantLimited {
 					st.Errorf("got limited=%v want %v", limited, tc.WantLimited)
 				}
@@ -184,12 +197,7 @@ func TestGetEventsInRangeWithTopologyToken(t *testing.T) {
 		events := r.Events()
 		_ = MustWriteEvents(t, db, events)
 
-		snapshot, err := db.NewDatabaseSnapshot(ctx)
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
-		if err != nil {
-			t.Fatal(err)
-		}
-		defer snapshot.Rollback() // nolint:errcheck
-
 			from, err := snapshot.MaxTopologicalPosition(ctx, r.ID)
 			if err != nil {
 				t.Fatalf("failed to get MaxTopologicalPosition: %s", err)
@@ -207,6 +215,7 @@ func TestGetEventsInRangeWithTopologyToken(t *testing.T) {
 			gots := snapshot.StreamEventsToEvents(nil, paginatedEvents)
 			test.AssertEventsEqual(t, gots, test.Reversed(events[len(events)-5:]))
 		})
+	})
 }
 
 /*
@@ -426,11 +435,8 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 		defer closeBase()
 		// At this point there should be no messages. We haven't sent anything
 		// yet.
-		snapshot, err := db.NewDatabaseSnapshot(ctx)
+
-		if err != nil {
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
-			t.Fatal(err)
-		}
-		defer snapshot.Rollback() // nolint:errcheck
 			_, events, err := snapshot.SendToDeviceUpdatesForSync(ctx, alice.ID, deviceID, 0, 100)
 			if err != nil {
 				t.Fatal(err)
@@ -438,6 +444,7 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 			if len(events) != 0 {
 				t.Fatal("first call should have no updates")
 			}
+		})
 
 		// Try sending a message.
 		streamPos, err := db.StoreNewSendForDeviceMessage(ctx, alice.ID, deviceID, gomatrixserverlib.SendToDeviceEvent{
@@ -449,9 +456,11 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 			t.Fatal(err)
 		}
 
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
 			// At this point we should get exactly one message. We're sending the sync position
 			// that we were given from the update and the send-to-device update will be updated
 			// in the database to reflect that this was the sync position we sent the message at.
+			var events []types.SendToDeviceEvent
 			streamPos, events, err = snapshot.SendToDeviceUpdatesForSync(ctx, alice.ID, deviceID, 0, streamPos)
 			if err != nil {
 				t.Fatal(err)
@@ -470,13 +479,17 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 			if len(events) != 1 {
 				t.Fatal("third call should have one update still")
 			}
+		})
+
 		err = db.CleanSendToDeviceUpdates(context.Background(), alice.ID, deviceID, streamPos)
 		if err != nil {
 			return
 		}
 
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
 			// At this point we should now have no updates, because we've progressed the sync
 			// position. Therefore the update from before will not be sent again.
+			var events []types.SendToDeviceEvent
 			_, events, err = snapshot.SendToDeviceUpdatesForSync(ctx, alice.ID, deviceID, streamPos, streamPos+10)
 			if err != nil {
 				t.Fatal(err)
@@ -494,6 +507,7 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 			if len(events) != 0 {
 				t.Fatal("fifth call should have no updates")
 			}
+		})
 
 		// Send some more messages and verify the ordering is correct ("in order of arrival")
 		var lastPos types.StreamPosition = 0
@@ -509,7 +523,8 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 			lastPos = streamPos
 		}
 
-		_, events, err = snapshot.SendToDeviceUpdatesForSync(ctx, alice.ID, deviceID, 0, lastPos)
+		WithSnapshot(t, db, func(snapshot storage.DatabaseSnapshot) {
+			_, events, err := snapshot.SendToDeviceUpdatesForSync(ctx, alice.ID, deviceID, 0, lastPos)
 			if err != nil {
 				t.Fatalf("unable to get events: %v", err)
 			}
@@ -522,6 +537,7 @@ func TestSendToDeviceBehaviour(t *testing.T) {
 				}
 			}
 		})
+	})
 }
 
 /*