Warp/server/internal/storage/push_test.go

package storage_test

import (
	"database/sql"
	"os"
	"slices"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// migratedPushDB opens a migrated database with empty push tables.
func migratedPushDB(t *testing.T, driver string) storage.Database {
	t.Helper()
	db := openTestDB(t, driver)
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	if driver == storage.DriverPostgres {
		pool, err := sql.Open("pgx", os.Getenv("WARP_TEST_POSTGRES_DSN"))
		if err != nil {
			t.Fatalf("open postgres to reset the push tables: %v", err)
		}
		defer func() { _ = pool.Close() }()
		if _, err := pool.ExecContext(t.Context(),
			`TRUNCATE push_pushers, push_notifications, push_room_counts`); err != nil {
			t.Fatalf("reset the push tables: %v", err)
		}
		if _, err := pool.ExecContext(t.Context(),
			`UPDATE push_sequences SET value = 0`); err != nil {
			t.Fatalf("reset the push counter: %v", err)
		}
	}
	return db
}

func TestPushers(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			push := migratedPushDB(t, driver).Push()

			pusher := func(localpart, device, key string) storage.Pusher {
				return storage.Pusher{
					Localpart: localpart, DeviceID: device, Kind: "http", AppID: "app", PushKey: key,
					AppDisplayName: "App", DeviceDisplayName: "Phone", Lang: "en",
					Data: []byte(`{"url":"https://push.example/_matrix/push/v1/notify"}`), PushKeyTSMS: 1,
				}
			}
			if err := push.SetPusher(ctx, pusher("alice", "D1", "k1"), false); err != nil {
				t.Fatal(err)
			}
			if err := push.SetPusher(ctx, pusher("alice", "D2", "k2"), false); err != nil {
				t.Fatal(err)
			}

			// Replacing keeps one row and takes the new values.
			updated := pusher("alice", "D1", "k1")
			updated.Lang = "fr"
			if err := push.SetPusher(ctx, updated, false); err != nil {
				t.Fatal(err)
			}
			got, err := push.Pushers(ctx, "alice")
			if err != nil || len(got) != 2 || got[0].Lang != "fr" || string(got[0].Data) != string(updated.Data) {
				t.Fatalf("Pushers = %+v, %v", got, err)
			}

			// The same key for another user takes it over, unless appending.
			if err := push.SetPusher(ctx, pusher("bob", "B1", "k1"), true); err != nil {
				t.Fatal(err)
			}
			if got, _ := push.Pushers(ctx, "alice"); len(got) != 2 {
				t.Errorf("appending removed alice's pusher: %+v", got)
			}
			if err := push.SetPusher(ctx, pusher("bob", "B1", "k1"), false); err != nil {
				t.Fatal(err)
			}
			if got, _ := push.Pushers(ctx, "alice"); len(got) != 1 || got[0].PushKey != "k2" {
				t.Errorf("taking the key over left alice with %+v", got)
			}

			if err := push.DeletePushersForDevices(ctx, "alice", []string{"D2"}); err != nil {
				t.Fatal(err)
			}
			if got, _ := push.Pushers(ctx, "alice"); len(got) != 0 {
				t.Errorf("after deleting the device: %+v", got)
			}
			if err := push.DeletePusher(ctx, "bob", "app", "k1"); err != nil {
				t.Fatal(err)
			}
			if err := push.DeletePusher(ctx, "bob", "app", "k1"); err != nil {
				t.Errorf("deleting a pusher twice: %v", err)
			}
			if got, _ := push.Pushers(ctx, "bob"); len(got) != 0 {
				t.Errorf("bob's pushers after the delete: %+v", got)
			}
		})
	}
}

func TestNotifications(t *testing.T) {
	const alice, bob = "@alice:warp.example", "@bob:warp.example"
	const room1, room2 = "!one:warp.example", "!two:warp.example"

	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			db := migratedPushDB(t, driver)
			push := db.Push()

			start, err := db.Sync().StreamPositions(ctx)
			if err != nil {
				t.Fatal(err)
			}

			note := func(user, room, event string, pos int64, highlight bool) storage.NewNotification {
				return storage.NewNotification{UserID: user, RoomID: room, EventID: event, EventPos: pos,
					TimestampMS: 1000 + pos, Highlight: highlight, Actions: []byte(`["notify"]`)}
			}
			last, err := push.AddNotifications(ctx, []storage.NewNotification{
				note(alice, room1, "$a", 10, false),
				note(alice, room1, "$b", 11, true),
				note(alice, room2, "$c", 12, false),
				note(bob, room1, "$b", 11, false),
			})
			if err != nil || last != start.Notification+4 {
				t.Fatalf("AddNotifications = %d, %v", last, err)
			}
			if positions, _ := db.Sync().StreamPositions(ctx); positions.Notification != last {
				t.Errorf("the stream is at %d, want %d", positions.Notification, last)
			}

			counts, err := push.NotificationCounts(ctx, alice, []string{room1, room2, "!none:x"})
			if err != nil {
				t.Fatal(err)
			}
			if counts[room1] != (storage.NotificationCounts{Notifications: 2, Highlights: 1}) ||
				counts[room2] != (storage.NotificationCounts{Notifications: 1}) ||
				counts["!none:x"] != (storage.NotificationCounts{}) {
				t.Errorf("counts = %+v", counts)
			}
			if total, _ := push.UnreadNotifications(ctx, alice); total != 3 {
				t.Errorf("UnreadNotifications = %d", total)
			}

			changed, err := push.NotificationCountChanges(ctx, alice, start.Notification, last)
			if err != nil || !slices.Equal(changed, []string{room1, room2}) {
				t.Errorf("NotificationCountChanges = %v, %v", changed, err)
			}

			// Reading up to the first event clears it alone, and moves the
			// stream for that room only.
			pos, moved, err := push.MarkNotificationsRead(ctx, alice, room1, storage.AllThreads, 10)
			if err != nil || !moved || pos != last+1 {
				t.Fatalf("MarkNotificationsRead = %d, %v, %v", pos, moved, err)
			}
			counts, _ = push.NotificationCounts(ctx, alice, []string{room1})
			if counts[room1] != (storage.NotificationCounts{Notifications: 1, Highlights: 1}) {
				t.Errorf("after reading one: %+v", counts)
			}
			if changed, _ := push.NotificationCountChanges(ctx, alice, last, pos); !slices.Equal(changed, []string{room1}) {
				t.Errorf("changes after the read = %v", changed)
			}

			// Reading the same again moves nothing.
			if again, moved, err := push.MarkNotificationsRead(ctx, alice, room1, storage.AllThreads, 10); err != nil || moved || again != pos {
				t.Errorf("reading again = %d, %v, %v", again, moved, err)
			}

			if _, _, err := push.MarkNotificationsRead(ctx, alice, room1, storage.AllThreads, 100); err != nil {
				t.Fatal(err)
			}
			counts, _ = push.NotificationCounts(ctx, alice, []string{room1})
			if counts[room1] != (storage.NotificationCounts{}) {
				t.Errorf("after reading everything: %+v", counts)
			}
			// Bob's are his own.
			if counts, _ := push.NotificationCounts(ctx, bob, []string{room1}); counts[room1].Notifications != 1 {
				t.Errorf("bob's counts = %+v", counts)
			}

			all, err := push.Notifications(ctx, alice, 0, 10, false)
			if err != nil || len(all) != 3 || all[0].EventID != "$c" || all[0].Read || !all[2].Read ||
				string(all[1].Actions) != `["notify"]` || !all[1].Highlight {
				t.Fatalf("Notifications = %+v, %v", all, err)
			}
			page, _ := push.Notifications(ctx, alice, all[0].ID, 1, false)
			if len(page) != 1 || page[0].EventID != "$b" {
				t.Errorf("the second page = %+v", page)
			}
			highlights, _ := push.Notifications(ctx, alice, 0, 10, true)
			if len(highlights) != 1 || highlights[0].EventID != "$b" {
				t.Errorf("highlights = %+v", highlights)
			}

			// Pruning takes only read notifications older than the cutoff.
			n, err := push.PruneNotifications(ctx, 1011)
			if err != nil || n != 1 {
				t.Errorf("PruneNotifications = %d, %v", n, err)
			}
			if rest, _ := push.Notifications(ctx, alice, 0, 10, false); len(rest) != 2 {
				t.Errorf("after pruning: %+v", rest)
			}
		})
	}
}

// TestThreadNotifications covers notifications filed by thread: a receipt
// for one thread, for the main timeline, and for the whole room each read
// only their own part, while the room totals stay the sum of every part.
func TestThreadNotifications(t *testing.T) {
	const alice = "@alice:warp.example"
	const room, other = "!one:warp.example", "!two:warp.example"
	const t1, t2 = "$root1", "$root2"

	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			push := migratedPushDB(t, driver).Push()

			note := func(event, thread string, pos int64, highlight bool) storage.NewNotification {
				return storage.NewNotification{UserID: alice, RoomID: room, EventID: event, EventPos: pos,
					ThreadID: thread, TimestampMS: 1000 + pos, Highlight: highlight, Actions: []byte(`["notify"]`)}
			}
			if _, err := push.AddNotifications(ctx, []storage.NewNotification{
				note(t1, storage.ThreadMain, 10, false), // a thread root is on the main timeline
				note("$a", t1, 11, true),
				note("$b", t1, 12, false),
				note("$c", t2, 13, false),
				note("$d", storage.ThreadMain, 14, false),
			}); err != nil {
				t.Fatal(err)
			}

			byThread := func() map[string]storage.NotificationCounts {
				t.Helper()
				counts, err := push.ThreadNotificationCounts(ctx, alice, []string{room, other})
				if err != nil {
					t.Fatal(err)
				}
				if _, ok := counts[other]; ok {
					t.Errorf("a room with nothing unread is in the result: %+v", counts[other])
				}
				return counts[room]
			}
			total := func() storage.NotificationCounts {
				t.Helper()
				counts, err := push.NotificationCounts(ctx, alice, []string{room})
				if err != nil {
					t.Fatal(err)
				}
				return counts[room]
			}

			got := byThread()
			if len(got) != 3 ||
				got[storage.ThreadMain] != (storage.NotificationCounts{Notifications: 2}) ||
				got[t1] != (storage.NotificationCounts{Notifications: 2, Highlights: 1}) ||
				got[t2] != (storage.NotificationCounts{Notifications: 1}) {
				t.Errorf("ThreadNotificationCounts = %+v", got)
			}
			if all, _ := push.Notifications(ctx, alice, 0, 10, false); len(all) != 5 || all[1].ThreadID != t2 {
				t.Errorf("Notifications = %+v", all)
			}

			// A receipt in thread 1 reads that thread, past everything else
			// in the room, and leaves the rest unread.
			if _, moved, err := push.MarkNotificationsRead(ctx, alice, room, t1, 100); err != nil || !moved {
				t.Fatalf("reading thread 1 = %v, %v", moved, err)
			}
			got = byThread()
			if _, ok := got[t1]; ok || len(got) != 2 {
				t.Errorf("after reading thread 1: %+v", got)
			}
			if c := total(); c != (storage.NotificationCounts{Notifications: 3}) {
				t.Errorf("the room total after reading thread 1 = %+v", c)
			}

			// A main receipt reads the main timeline up to its event.
			if _, moved, err := push.MarkNotificationsRead(ctx, alice, room, storage.ThreadMain, 10); err != nil || !moved {
				t.Fatalf("reading the main timeline = %v, %v", moved, err)
			}
			got = byThread()
			if got[storage.ThreadMain] != (storage.NotificationCounts{Notifications: 1}) ||
				got[t2] != (storage.NotificationCounts{Notifications: 1}) {
				t.Errorf("after the main receipt: %+v", got)
			}
			// A receipt for a thread with nothing unread moves nothing.
			if _, moved, err := push.MarkNotificationsRead(ctx, alice, room, t1, 100); err != nil || moved {
				t.Errorf("reading thread 1 again = %v, %v", moved, err)
			}

			// An unthreaded receipt reads everything.
			if _, _, err := push.MarkNotificationsRead(ctx, alice, room, storage.AllThreads, 100); err != nil {
				t.Fatal(err)
			}
			if got := byThread(); len(got) != 0 {
				t.Errorf("after the unthreaded receipt: %+v", got)
			}
			if c := total(); c != (storage.NotificationCounts{}) {
				t.Errorf("the room total after the unthreaded receipt = %+v", c)
			}
		})
	}
}