From 2ea06dacd804dea3889333b3b8d94eb3eb96e28c Mon Sep 17 00:00:00 2001
From: Kegan Dougal <kegan@matrix.org>
Date: Tue, 28 Jul 2020 15:06:14 +0100
Subject: [PATCH] Add tests and fix maxOffset bug

---
 .../storage/postgres/key_changes_table.go     |  5 +-
 .../storage/sqlite3/key_changes_table.go      |  5 +-
 keyserver/storage/storage_test.go             | 57 +++++++++++++++++++
 3 files changed, 61 insertions(+), 6 deletions(-)
 create mode 100644 keyserver/storage/storage_test.go

diff --git a/keyserver/storage/postgres/key_changes_table.go b/keyserver/storage/postgres/key_changes_table.go
index 246131b92..dabdafcb6 100644
--- a/keyserver/storage/postgres/key_changes_table.go
+++ b/keyserver/storage/postgres/key_changes_table.go
@@ -83,15 +83,14 @@ func (s *keyChangesStatements) SelectKeyChanges(
 		return nil, 0, err
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "selectKeyChangesStmt: rows.close() failed")
-	var maxOffset int64
 	for rows.Next() {
 		var userID string
 		var offset int64
 		if err := rows.Scan(&userID, &offset); err != nil {
 			return nil, 0, err
 		}
-		if offset > maxOffset {
-			maxOffset = offset
+		if offset > latestOffset {
+			latestOffset = offset
 		}
 		userIDs = append(userIDs, userID)
 	}
diff --git a/keyserver/storage/sqlite3/key_changes_table.go b/keyserver/storage/sqlite3/key_changes_table.go
index 3264e9e61..b830214d1 100644
--- a/keyserver/storage/sqlite3/key_changes_table.go
+++ b/keyserver/storage/sqlite3/key_changes_table.go
@@ -83,15 +83,14 @@ func (s *keyChangesStatements) SelectKeyChanges(
 		return nil, 0, err
 	}
 	defer internal.CloseAndLogIfError(ctx, rows, "selectKeyChangesStmt: rows.close() failed")
-	var maxOffset int64
 	for rows.Next() {
 		var userID string
 		var offset int64
 		if err := rows.Scan(&userID, &offset); err != nil {
 			return nil, 0, err
 		}
-		if offset > maxOffset {
-			maxOffset = offset
+		if offset > latestOffset {
+			latestOffset = offset
 		}
 		userIDs = append(userIDs, userID)
 	}
diff --git a/keyserver/storage/storage_test.go b/keyserver/storage/storage_test.go
new file mode 100644
index 000000000..889724789
--- /dev/null
+++ b/keyserver/storage/storage_test.go
@@ -0,0 +1,57 @@
+package storage
+
+import (
+	"context"
+	"reflect"
+	"testing"
+)
+
+var ctx = context.Background()
+
+func MustNotError(t *testing.T, err error) {
+	t.Helper()
+	if err == nil {
+		return
+	}
+	t.Fatalf("operation failed: %s", err)
+}
+
+func TestKeyChanges(t *testing.T) {
+	db, err := NewDatabase("file::memory:", nil)
+	if err != nil {
+		t.Fatalf("Failed to NewDatabase: %s", err)
+	}
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 0, "@alice:localhost"))
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 1, "@bob:localhost"))
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 2, "@charlie:localhost"))
+	userIDs, latest, err := db.KeyChanges(ctx, 0, 1)
+	if err != nil {
+		t.Fatalf("Failed to KeyChanges: %s", err)
+	}
+	if latest != 2 {
+		t.Fatalf("KeyChanges: got latest=%d want 2", latest)
+	}
+	if !reflect.DeepEqual(userIDs, []string{"@charlie:localhost"}) {
+		t.Fatalf("KeyChanges: wrong user_ids: %v", userIDs)
+	}
+}
+
+func TestKeyChangesNoDupes(t *testing.T) {
+	db, err := NewDatabase("file::memory:", nil)
+	if err != nil {
+		t.Fatalf("Failed to NewDatabase: %s", err)
+	}
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 0, "@alice:localhost"))
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 1, "@alice:localhost"))
+	MustNotError(t, db.StoreKeyChange(ctx, 0, 2, "@alice:localhost"))
+	userIDs, latest, err := db.KeyChanges(ctx, 0, 0)
+	if err != nil {
+		t.Fatalf("Failed to KeyChanges: %s", err)
+	}
+	if latest != 2 {
+		t.Fatalf("KeyChanges: got latest=%d want 2", latest)
+	}
+	if !reflect.DeepEqual(userIDs, []string{"@alice:localhost"}) {
+		t.Fatalf("KeyChanges: wrong user_ids: %v", userIDs)
+	}
+}