Add EventStateKeysTable tests

roomserver/storage/postgres/event_state_keys_table.go

@@ -76,12 +76,12 @@ type eventStateKeyStatements struct {
 	bulkSelectEventStateKeyStmt    *sql.Stmt
 }
 
-func createEventStateKeysTable(db *sql.DB) error {
+func CreateEventStateKeysTable(db *sql.DB) error {
 	_, err := db.Exec(eventStateKeysSchema)
 	return err
 }
 
-func prepareEventStateKeysTable(db *sql.DB) (tables.EventStateKeys, error) {
+func PrepareEventStateKeysTable(db *sql.DB) (tables.EventStateKeys, error) {
 	s := &eventStateKeyStatements{}
 
 	return s, sqlutil.StatementList{
@@ -123,9 +123,9 @@ func (s *eventStateKeyStatements) BulkSelectEventStateKeyNID(
 	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectEventStateKeyNID: rows.close() failed")
 
 	result := make(map[string]types.EventStateKeyNID, len(eventStateKeys))
+	var stateKey string
+	var stateKeyNID int64
 	for rows.Next() {
-		var stateKey string
-		var stateKeyNID int64
 		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
 			return nil, err
 		}
@@ -149,9 +149,9 @@ func (s *eventStateKeyStatements) BulkSelectEventStateKey(
 	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectEventStateKey: rows.close() failed")
 
 	result := make(map[types.EventStateKeyNID]string, len(eventStateKeyNIDs))
+	var stateKey string
+	var stateKeyNID int64
 	for rows.Next() {
-		var stateKey string
-		var stateKeyNID int64
 		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
 			return nil, err
 		}

roomserver/storage/postgres/storage.go

@@ -68,7 +68,7 @@ func Open(base *base.BaseDendrite, dbProperties *config.DatabaseOptions, cache c
 }
 
 func (d *Database) create(db *sql.DB) error {
-	if err := createEventStateKeysTable(db); err != nil {
+	if err := CreateEventStateKeysTable(db); err != nil {
 		return err
 	}
 	if err := createEventTypesTable(db); err != nil {
@@ -112,7 +112,7 @@ func (d *Database) create(db *sql.DB) error {
 }
 
 func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.RoomServerCaches) error {
-	eventStateKeys, err := prepareEventStateKeysTable(db)
+	eventStateKeys, err := PrepareEventStateKeysTable(db)
 	if err != nil {
 		return err
 	}

roomserver/storage/sqlite3/event_state_keys_table.go

@@ -71,12 +71,12 @@ type eventStateKeyStatements struct {
 	bulkSelectEventStateKeyStmt    *sql.Stmt
 }
 
-func createEventStateKeysTable(db *sql.DB) error {
+func CreateEventStateKeysTable(db *sql.DB) error {
 	_, err := db.Exec(eventStateKeysSchema)
 	return err
 }
 
-func prepareEventStateKeysTable(db *sql.DB) (tables.EventStateKeys, error) {
+func PrepareEventStateKeysTable(db *sql.DB) (tables.EventStateKeys, error) {
 	s := &eventStateKeyStatements{
 		db: db,
 	}
@@ -128,9 +128,9 @@ func (s *eventStateKeyStatements) BulkSelectEventStateKeyNID(
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectEventStateKeyNID: rows.close() failed")
 	result := make(map[string]types.EventStateKeyNID, len(eventStateKeys))
+	var stateKey string
+	var stateKeyNID int64
 	for rows.Next() {
-		var stateKey string
-		var stateKeyNID int64
 		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
 			return nil, err
 		}
@@ -159,9 +159,9 @@ func (s *eventStateKeyStatements) BulkSelectEventStateKey(
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectEventStateKey: rows.close() failed")
 	result := make(map[types.EventStateKeyNID]string, len(eventStateKeyNIDs))
+	var stateKey string
+	var stateKeyNID int64
 	for rows.Next() {
-		var stateKey string
-		var stateKeyNID int64
 		if err := rows.Scan(&stateKey, &stateKeyNID); err != nil {
 			return nil, err
 		}

roomserver/storage/sqlite3/storage.go

@@ -77,7 +77,7 @@ func Open(base *base.BaseDendrite, dbProperties *config.DatabaseOptions, cache c
 }
 
 func (d *Database) create(db *sql.DB) error {
-	if err := createEventStateKeysTable(db); err != nil {
+	if err := CreateEventStateKeysTable(db); err != nil {
 		return err
 	}
 	if err := createEventTypesTable(db); err != nil {
@@ -121,7 +121,7 @@ func (d *Database) create(db *sql.DB) error {
 }
 
 func (d *Database) prepare(db *sql.DB, writer sqlutil.Writer, cache caching.RoomServerCaches) error {
-	eventStateKeys, err := prepareEventStateKeysTable(db)
+	eventStateKeys, err := PrepareEventStateKeysTable(db)
 	if err != nil {
 		return err
 	}

roomserver/storage/tables/event_json_table_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/matrix-org/dendrite/test"
 )
 
-func mustCreateTable(t *testing.T, dbType test.DBType) (tables.EventJSON, func()) {
+func mustCreateEventJSONTable(t *testing.T, dbType test.DBType) (tables.EventJSON, func()) {
 	t.Helper()
 	connStr, close := test.PrepareDBConnectionString(t, dbType)
 	db, err := sqlutil.Open(&config.DatabaseOptions{
@@ -48,7 +48,7 @@ func mustCreateTable(t *testing.T, dbType test.DBType) (tables.EventJSON, func()
 
 func Test_EventJSONTable(t *testing.T) {
 	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
-		tab, close := mustCreateTable(t, dbType)
+		tab, close := mustCreateEventJSONTable(t, dbType)
 		defer close()
 		// create some dummy data
 		for i := 0; i < 10; i++ {

roomserver/storage/tables/event_state_keys_test.go

@@ -0,0 +1,94 @@
+package tables_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/postgres"
+	"github.com/matrix-org/dendrite/roomserver/storage/sqlite3"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+	"github.com/matrix-org/dendrite/setup/config"
+	"github.com/matrix-org/dendrite/test"
+)
+
+func mustCreateEventStateKeysTable(t *testing.T, dbType test.DBType) (tables.EventStateKeys, func()) {
+	t.Helper()
+	connStr, close := test.PrepareDBConnectionString(t, dbType)
+	db, err := sqlutil.Open(&config.DatabaseOptions{
+		ConnectionString: config.DataSource(connStr),
+	}, sqlutil.NewExclusiveWriter())
+	if err != nil {
+		t.Fatalf("failed to open db: %s", err)
+	}
+	var tab tables.EventStateKeys
+	switch dbType {
+	case test.DBTypePostgres:
+		err = postgres.CreateEventStateKeysTable(db)
+		if err != nil {
+			t.Fatalf("failed to create EventJSON table: %s", err)
+		}
+		tab, err = postgres.PrepareEventStateKeysTable(db)
+	case test.DBTypeSQLite:
+		err = sqlite3.CreateEventStateKeysTable(db)
+		if err != nil {
+			t.Fatalf("failed to create EventJSON table: %s", err)
+		}
+		tab, err = sqlite3.PrepareEventStateKeysTable(db)
+	}
+	if err != nil {
+		t.Fatalf("failed to create table: %s", err)
+	}
+
+	return tab, close
+}
+
+func Test_EventStateKeysTable(t *testing.T) {
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		tab, close := mustCreateEventStateKeysTable(t, dbType)
+		defer close()
+		ctx := context.Background()
+		var stateKeyNID, gotEventStateKey types.EventStateKeyNID
+		var err error
+		// create some dummy data
+		for i := 0; i < 10; i++ {
+			stateKey := fmt.Sprintf("@user%d:localhost", i)
+			if stateKeyNID, err = tab.InsertEventStateKeyNID(
+				ctx, nil, stateKey,
+			); err != nil {
+				t.Fatalf("unable to insert eventJSON: %s", err)
+			}
+			gotEventStateKey, err = tab.SelectEventStateKeyNID(ctx, nil, stateKey)
+			if err != nil {
+				t.Fatalf("failed to get eventStateKeyNID: %s", err)
+			}
+			if stateKeyNID != gotEventStateKey {
+				t.Fatalf("expected eventStateKey %d, but got %d", stateKeyNID, gotEventStateKey)
+			}
+		}
+		stateKeyNIDsMap, err := tab.BulkSelectEventStateKeyNID(ctx, nil, []string{"@user0:localhost", "@user1:localhost"})
+		if err != nil {
+			t.Fatalf("failed to get EventStateKeyNIDs: %s", err)
+		}
+		wantStateKeyNIDs := make([]types.EventStateKeyNID, 0, len(stateKeyNIDsMap))
+		for _, nid := range stateKeyNIDsMap {
+			wantStateKeyNIDs = append(wantStateKeyNIDs, nid)
+		}
+		stateKeyNIDs, err := tab.BulkSelectEventStateKey(ctx, nil, wantStateKeyNIDs)
+		if err != nil {
+			t.Fatalf("failed to get EventStateKeyNIDs: %s", err)
+		}
+		// verify that BulkSelectEventStateKeyNID and BulkSelectEventStateKey return the same values
+		for userID, nid := range stateKeyNIDsMap {
+			if v, ok := stateKeyNIDs[nid]; ok {
+				if v != userID {
+					t.Fatalf("userID does not match: %s != %s", userID, v)
+				}
+			} else {
+				t.Fatalf("unable to find %d in result set", nid)
+			}
+		}
+	})
+}