Add tests

src/github.com/matrix-org/dendrite/roomserver/query/query.go

@@ -27,9 +27,19 @@ import (
 	"github.com/matrix-org/util"
 )
 
+// RoomserverQueryAPIEventDB has a convenience API to fetch events directly by
+// EventIDs.
+type RoomserverQueryAPIEventDB interface {
+	// Look up the Events for a list of event IDs. Does not error if event was
+	// not found.
+	// Returns an error if the retrieval went wrong.
+	EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error)
+}
+
 // RoomserverQueryAPIDatabase has the storage APIs needed to implement the query API.
 type RoomserverQueryAPIDatabase interface {
 	state.RoomStateDatabase
+	RoomserverQueryAPIEventDB
 	// Look up the numeric ID for the room.
 	// Returns 0 if the room doesn't exists.
 	// Returns an error if there was a problem talking to the database.
@@ -459,7 +469,7 @@ func (r *RoomserverQueryAPI) QueryStateAndAuthChain(
 	}
 
 	response.StateEvents = stateEvents
-	response.AuthChainEvents, err = r.getAuthChain(ctx, request.AuthEventIDs)
+	response.AuthChainEvents, err = getAuthChain(ctx, r.DB, request.AuthEventIDs)
 	return err
 }
 
@@ -468,8 +478,8 @@ func (r *RoomserverQueryAPI) QueryStateAndAuthChain(
 // auth_events section, and all their auth_events, recursively.
 // The returned set of events contain the given events.
 // Will *not* error if we don't have all auth events.
-func (r *RoomserverQueryAPI) getAuthChain(
-	ctx context.Context, authEventIDs []string,
+func getAuthChain(
+	ctx context.Context, dB RoomserverQueryAPIEventDB, authEventIDs []string,
 ) ([]gomatrixserverlib.Event, error) {
 	var authEvents []gomatrixserverlib.Event
 
@@ -486,17 +496,7 @@ func (r *RoomserverQueryAPI) getAuthChain(
 	// Check if there's anything left to do
 	for len(eventsToFetch) > 0 {
 		// Convert eventIDs to events. First need to fetch NIDs
-		nidMap, err := r.DB.EventNIDs(ctx, eventsToFetch)
-		if err != nil {
-			return nil, err
-		}
-
-		var nids []types.EventNID
-		for _, nid := range nidMap {
-			nids = append(nids, nid)
-		}
-
-		events, err := r.DB.Events(ctx, nids)
+		events, err := dB.EventsFromIDs(ctx, eventsToFetch)
 		if err != nil {
 			return nil, err
 		}

src/github.com/matrix-org/dendrite/roomserver/query/query_test.go

@@ -0,0 +1,174 @@
+// Copyright 2017 Vector Creations Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package query
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"sort"
+
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/gomatrixserverlib"
+)
+
+// used to implement RoomserverQueryAPIEventDB to test getAuthChain
+type getEventDB struct {
+	eventMap map[string]gomatrixserverlib.Event
+}
+
+func createEventDB() *getEventDB {
+	return &getEventDB{
+		eventMap: make(map[string]gomatrixserverlib.Event),
+	}
+}
+
+// Adds a fake event to the storage with given auth events.
+func (db *getEventDB) addFakeEvent(eventID string, authIDs []string) error {
+	authEvents := []gomatrixserverlib.EventReference{}
+	for _, authID := range authIDs {
+		authEvents = append(authEvents, gomatrixserverlib.EventReference{
+			EventID: authID,
+		})
+	}
+
+	builder := map[string]interface{}{
+		"event_id":    eventID,
+		"auth_events": authEvents,
+	}
+
+	eventJSON, err := json.Marshal(&builder)
+	if err != nil {
+		return err
+	}
+
+	event, err := gomatrixserverlib.NewEventFromTrustedJSON(eventJSON, false)
+	if err != nil {
+		return err
+	}
+
+	db.eventMap[eventID] = event
+
+	return nil
+}
+
+// Adds multiple events at once, each entry in the map is an eventID and set of
+// auth events that are converted to an event and added.
+func (db *getEventDB) addFakeEvents(graph map[string][]string) error {
+	for eventID, authIDs := range graph {
+		err := db.addFakeEvent(eventID, authIDs)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// EventsFromIDs implements RoomserverQueryAPIEventDB
+func (db *getEventDB) EventsFromIDs(ctx context.Context, eventIDs []string) (res []types.Event, err error) {
+	for _, evID := range eventIDs {
+		res = append(res, types.Event{
+			EventNID: 0,
+			Event:    db.eventMap[evID],
+		})
+	}
+
+	return
+}
+
+// Returns if the slices are equal after sorting them.
+func compareUnsortedStringSlices(a []string, b []string) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	sort.Strings(a)
+	sort.Strings(b)
+
+	for i := range a {
+		if a[i] != b[i] {
+			return false
+		}
+	}
+
+	return true
+}
+
+func TestGetAuthChainSingle(t *testing.T) {
+	db := createEventDB()
+
+	err := db.addFakeEvents(map[string][]string{
+		"a": {},
+		"b": {"a"},
+		"c": {"a", "b"},
+		"d": {"b", "c"},
+		"e": {"a", "d"},
+	})
+
+	if err != nil {
+		t.Fatalf("Failed to add events to db: %v", err)
+	}
+
+	result, err := getAuthChain(context.TODO(), db, []string{"e"})
+	if err != nil {
+		t.Fatalf("getAuthChain failed: %v", err)
+	}
+
+	var returnedIDs []string
+	for _, event := range result {
+		returnedIDs = append(returnedIDs, event.EventID())
+	}
+
+	expectedIDs := []string{"a", "b", "c", "d", "e"}
+
+	if !compareUnsortedStringSlices(expectedIDs, returnedIDs) {
+		t.Fatalf("returnedIDs got '%v', expected '%v'", returnedIDs, expectedIDs)
+	}
+}
+
+func TestGetAuthChainMultiple(t *testing.T) {
+	db := createEventDB()
+
+	err := db.addFakeEvents(map[string][]string{
+		"a": {},
+		"b": {"a"},
+		"c": {"a", "b"},
+		"d": {"b", "c"},
+		"e": {"a", "d"},
+		"f": {"a", "b", "c"},
+	})
+
+	if err != nil {
+		t.Fatalf("Failed to add events to db: %v", err)
+	}
+
+	result, err := getAuthChain(context.TODO(), db, []string{"e", "f"})
+	if err != nil {
+		t.Fatalf("getAuthChain failed: %v", err)
+	}
+
+	var returnedIDs []string
+	for _, event := range result {
+		returnedIDs = append(returnedIDs, event.EventID())
+	}
+
+	expectedIDs := []string{"a", "b", "c", "d", "e", "f"}
+
+	if !compareUnsortedStringSlices(expectedIDs, returnedIDs) {
+		t.Fatalf("returnedIDs got '%v', expected '%v'", returnedIDs, expectedIDs)
+	}
+}

src/github.com/matrix-org/dendrite/roomserver/storage/storage.go

@@ -651,6 +651,21 @@ func (d *Database) GetMembershipEventNIDsForRoom(
 	return d.statements.selectMembershipsFromRoom(ctx, roomNID)
 }
 
+// EventsFromIDs implements query.RoomserverQueryAPIEventDB
+func (d *Database) EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error) {
+	nidMap, err := d.EventNIDs(ctx, eventIDs)
+	if err != nil {
+		return nil, err
+	}
+
+	var nids []types.EventNID
+	for _, nid := range nidMap {
+		nids = append(nids, nid)
+	}
+
+	return d.Events(ctx, nids)
+}
+
 type transaction struct {
 	ctx context.Context
 	txn *sql.Tx