Warp/server/internal/storage/relations_test.go

package storage_test

import (
	"slices"
	"testing"

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

// writeMessage appends a timeline event with the given projections and
// returns its stream ordering.
func writeMessage(t *testing.T, rooms storage.RoomStore, ev storage.StoredEvent) int64 {
	t.Helper()
	pos, err := rooms.Write(t.Context(), storage.RoomWrite{Event: ev, Extremities: []string{ev.EventID}})
	if err != nil {
		t.Fatalf("Write %s: %v", ev.EventID, err)
	}
	return pos
}

func relationEvent(id, sender, parent, relType string) storage.StoredEvent {
	e := newTestEvent(testRoomID, id, "m.room.message", nil)
	e.Sender = sender
	e.RelatesTo, e.RelType = parent, relType
	return e
}

func TestRelationsAndAggregates(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedRoomDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}

			writeMessage(t, rooms, newTestEvent(testRoomID, "$root", "m.room.message", nil))
			writeMessage(t, rooms, relationEvent("$t1", "@bob:warp.example", "$root", storage.RelThread))
			last := writeMessage(t, rooms, relationEvent("$t2", "@carol:warp.example", "$root", storage.RelThread))
			writeMessage(t, rooms, relationEvent("$edit", "@alice:warp.example", "$root", storage.RelReplace))
			writeMessage(t, rooms, relationEvent("$ref", "@bob:warp.example", "$root", storage.RelReference))
			writeMessage(t, rooms, relationEvent("$t1edit", "@bob:warp.example", "$t1", storage.RelReplace))

			children, err := rooms.Relations(ctx, storage.RelationsRequest{
				RoomID: testRoomID, ParentIDs: []string{"$root"},
				Direction: storage.DirectionBackward, Limit: 10,
			})
			if err != nil {
				t.Fatalf("Relations: %v", err)
			}
			if got := eventIDs(children); !slices.Equal(got, []string{"$ref", "$edit", "$t2", "$t1"}) {
				t.Errorf("Relations = %v", got)
			}
			if children[0].RelatesTo != "$root" || children[0].RelType != storage.RelReference {
				t.Errorf("the projections did not round-trip: %+v", children[0])
			}

			threads, err := rooms.Relations(ctx, storage.RelationsRequest{
				RoomID: testRoomID, ParentIDs: []string{"$root", "$t1"}, RelType: storage.RelReplace,
				Direction: storage.DirectionForward, Limit: 10,
			})
			if err != nil {
				t.Fatalf("Relations by type: %v", err)
			}
			if got := eventIDs(threads); !slices.Equal(got, []string{"$edit", "$t1edit"}) {
				t.Errorf("m.replace children of two parents = %v", got)
			}

			ids, err := rooms.ChildEventIDs(ctx, testRoomID, []string{"$t1"}, "", "")
			if err != nil || !slices.Equal(ids, []string{"$t1edit"}) {
				t.Errorf("ChildEventIDs = %v, %v", ids, err)
			}

			aggregates, err := rooms.RelationAggregates(ctx, testRoomID, []string{"$root"}, 0)
			if err != nil {
				t.Fatalf("RelationAggregates: %v", err)
			}
			threadCount, latest := 0, ""
			for _, a := range aggregates {
				if a.RelType == storage.RelThread {
					threadCount += a.Count
					if a.LatestOrdering == last {
						latest = a.LatestEventID
					}
				}
			}
			if threadCount != 2 || latest != "$t2" || len(aggregates) != 4 {
				t.Errorf("aggregates = %+v", aggregates)
			}

			refs, err := rooms.RelationEventIDs(ctx, testRoomID, []string{"$root"}, storage.RelReference, 0, 10)
			if err != nil || !slices.Equal(refs["$root"], []string{"$ref"}) {
				t.Errorf("RelationEventIDs = %v, %v", refs, err)
			}

			// Bounded above, the later replies are not counted.
			bounded, err := rooms.RelationAggregates(ctx, testRoomID, []string{"$root"}, last-1)
			if err != nil {
				t.Fatalf("RelationAggregates bounded: %v", err)
			}
			for _, a := range bounded {
				if a.RelType == storage.RelThread && a.LatestEventID == "$t2" {
					t.Errorf("a bounded aggregate saw the event above its bound: %+v", a)
				}
			}
		})
	}
}

func TestThreadsListing(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedRoomDB(t, driver).Rooms()
			if err := rooms.CreateRoom(ctx, newTestRoom(testRoomID)); err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}
			writeMessage(t, rooms, newTestEvent(testRoomID, "$a", "m.room.message", nil))
			b := newTestEvent(testRoomID, "$b", "m.room.message", nil)
			b.Sender = "@bob:warp.example"
			writeMessage(t, rooms, b)
			writeMessage(t, rooms, relationEvent("$a1", "@alice:warp.example", "$a", storage.RelThread))
			writeMessage(t, rooms, relationEvent("$b1", "@alice:warp.example", "$b", storage.RelThread))
			// A thread whose root this server does not hold is not listed.
			writeMessage(t, rooms, relationEvent("$x1", "@alice:warp.example", "$missing", storage.RelThread))

			all, err := rooms.Threads(ctx, storage.ThreadsRequest{RoomID: testRoomID, Limit: 10})
			if err != nil {
				t.Fatalf("Threads: %v", err)
			}
			if got := threadRoots(all); !slices.Equal(got, []string{"$b", "$a"}) {
				t.Errorf("Threads = %v, want newest activity first", got)
			}

			page, err := rooms.Threads(ctx, storage.ThreadsRequest{RoomID: testRoomID, Limit: 1, Before: all[0].LatestOrdering})
			if err != nil || !slices.Equal(threadRoots(page), []string{"$a"}) {
				t.Errorf("Threads after the first = %v, %v", threadRoots(page), err)
			}

			bobs, err := rooms.Threads(ctx, storage.ThreadsRequest{
				RoomID: testRoomID, ParticipatedBy: "@bob:warp.example", Limit: 10,
			})
			if err != nil || !slices.Equal(threadRoots(bobs), []string{"$b"}) {
				t.Errorf("Threads bob took part in = %v, %v", threadRoots(bobs), err)
			}
		})
	}
}

func TestSearch(t *testing.T) {
	for _, driver := range testDrivers {
		t.Run(driver, func(t *testing.T) {
			ctx := t.Context()
			rooms := migratedRoomDB(t, driver).Rooms()
			const other = "!other:warp.example"
			for _, id := range []string{testRoomID, other} {
				if err := rooms.CreateRoom(ctx, newTestRoom(id)); err != nil {
					t.Fatalf("CreateRoom: %v", err)
				}
			}

			message := func(roomID, id, body string) int64 {
				e := newTestEvent(roomID, id, "m.room.message", nil)
				e.SearchKey, e.SearchText = "content.body", body
				return writeMessage(t, rooms, e)
			}
			message(testRoomID, "$hello", "Hello wonderful world")
			message(testRoomID, "$cafe", "Meet at the café tomorrow")
			second := message(testRoomID, "$hello2", "hello again, hello")
			message(other, "$elsewhere", "hello from elsewhere")
			topic := newTestEvent(testRoomID, "$topic", "m.room.topic", stateKey(""))
			topic.SearchKey, topic.SearchText = "content.topic", "a room to say hello in"
			writeMessage(t, rooms, topic)

			search := func(req storage.SearchRequest) ([]string, int) {
				t.Helper()
				if req.Limit == 0 {
					req.Limit = 10
				}
				if req.Rooms == nil {
					req.Rooms = []storage.SearchRange{{RoomID: testRoomID}}
				}
				results, total, err := rooms.Search(ctx, req)
				if err != nil {
					t.Fatalf("Search(%q): %v", req.Term, err)
				}
				out := make([]string, 0, len(results))
				for _, r := range results {
					out = append(out, r.Event.EventID)
				}
				return out, total
			}

			got, total := search(storage.SearchRequest{Term: "hello", Order: storage.SearchByRecent})
			if !slices.Equal(got, []string{"$topic", "$hello2", "$hello"}) || total != 3 {
				t.Errorf("recent search for hello = %v (total %d)", got, total)
			}
			// A prefix matches, and diacritics are folded on SQLite; the
			// tsvector 'simple' configuration keeps them, so the term is
			// spelled as stored.
			if got, _ := search(storage.SearchRequest{Term: "hel"}); len(got) != 3 {
				t.Errorf("prefix search = %v", got)
			}
			if got, _ := search(storage.SearchRequest{Term: "café"}); !slices.Equal(got, []string{"$cafe"}) {
				t.Errorf("search for café = %v", got)
			}
			// Operators in the term are words, not syntax.
			if got, _ := search(storage.SearchRequest{Term: `"hello" & (world`}); !slices.Equal(got, []string{"$hello"}) {
				t.Errorf("a term full of operators = %v, want the one event with both words", got)
			}
			if got, _ := search(storage.SearchRequest{Term: "!!!"}); len(got) != 0 {
				t.Errorf("a term with no words found %v", got)
			}

			if got, _ := search(storage.SearchRequest{Term: "hello", Keys: []string{"content.topic"}}); !slices.Equal(got, []string{"$topic"}) {
				t.Errorf("search by key = %v", got)
			}
			if got, _ := search(storage.SearchRequest{
				Term: "hello", Rooms: []storage.SearchRange{{RoomID: testRoomID, High: second - 1}},
			}); !slices.Equal(got, []string{"$hello"}) {
				t.Errorf("search bounded below the second hello = %v", got)
			}
			if got, total := search(storage.SearchRequest{
				Term: "hello", Rooms: []storage.SearchRange{{RoomID: testRoomID}, {RoomID: other}},
			}); len(got) != 4 || total != 4 {
				t.Errorf("search across two rooms = %v (total %d)", got, total)
			}
			if got, _ := search(storage.SearchRequest{
				Term: "hello", Filter: storage.TimelineFilter{NotTypes: []string{"m.room.topic"}}, Order: storage.SearchByRecent,
			}); !slices.Equal(got, []string{"$hello2", "$hello"}) {
				t.Errorf("filtered search = %v", got)
			}

			// Paging by recency, then by rank.
			page1, _ := search(storage.SearchRequest{Term: "hello", Order: storage.SearchByRecent, Limit: 2})
			if len(page1) != 2 {
				t.Fatalf("first page = %v", page1)
			}
			if page2, _ := search(storage.SearchRequest{Term: "hello", Order: storage.SearchByRank, Offset: 2, Limit: 2}); len(page2) != 1 {
				t.Errorf("ranked second page = %v", page2)
			}

			// A redacted event is no longer found by what it said.
			redaction := newTestEvent(testRoomID, "$redaction", "m.room.redaction", nil)
			if _, err := rooms.Write(ctx, storage.RoomWrite{
				Event: redaction, Extremities: []string{"$redaction"},
				Redacts: "$hello2", RedactedJSON: []byte(`{"type":"m.room.message"}`),
			}); err != nil {
				t.Fatalf("Write redaction: %v", err)
			}
			if got, _ := search(storage.SearchRequest{Term: "again"}); len(got) != 0 {
				t.Errorf("the redacted event is still found: %v", got)
			}

			if err := rooms.PurgeRoom(ctx, testRoomID); err != nil {
				t.Fatalf("PurgeRoom: %v", err)
			}
			if got, _ := search(storage.SearchRequest{Term: "hello"}); len(got) != 0 {
				t.Errorf("a purged room is still searched: %v", got)
			}
		})
	}
}

func eventIDs(events []storage.StoredEvent) []string {
	out := make([]string, 0, len(events))
	for _, e := range events {
		out = append(out, e.EventID)
	}
	return out
}

func threadRoots(threads []storage.Thread) []string {
	out := make([]string, 0, len(threads))
	for _, t := range threads {
		out = append(out, t.RootID)
	}
	return out
}