From ead825b1bf25de9e22658520c3d2d9c38ec7ec80 Mon Sep 17 00:00:00 2001
From: Devon Hudson <devonhudson@librem.one>
Date: Thu, 22 Dec 2022 10:38:34 -0700
Subject: [PATCH] Add fed storage tests for assumed offline & relay servers

---
 federationapi/storage/storage_test.go | 93 +++++++++++++++++++++++++--
 1 file changed, 89 insertions(+), 4 deletions(-)

diff --git a/federationapi/storage/storage_test.go b/federationapi/storage/storage_test.go
index 5b57d40d4..f05055227 100644
--- a/federationapi/storage/storage_test.go
+++ b/federationapi/storage/storage_test.go
@@ -6,14 +6,13 @@ import (
 	"testing"
 	"time"
 
-	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/util"
-	"github.com/stretchr/testify/assert"
-
 	"github.com/matrix-org/dendrite/federationapi/storage"
 	"github.com/matrix-org/dendrite/setup/config"
 	"github.com/matrix-org/dendrite/test"
 	"github.com/matrix-org/dendrite/test/testrig"
+	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/util"
+	"github.com/stretchr/testify/assert"
 )
 
 func mustCreateFederationDatabase(t *testing.T, dbType test.DBType) (storage.Database, func()) {
@@ -246,3 +245,89 @@ func TestInboundPeeking(t *testing.T) {
 		assert.ElementsMatch(t, gotPeekIDs, peekIDs)
 	})
 }
+
+func TestServersAssumedOffline(t *testing.T) {
+	server1 := gomatrixserverlib.ServerName("server1")
+	server2 := gomatrixserverlib.ServerName("server2")
+
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		db, closeDB := mustCreateFederationDatabase(t, dbType)
+		defer closeDB()
+
+		err := db.SetServerAssumedOffline(server1)
+		assert.Nil(t, err)
+
+		isOffline, err := db.IsServerAssumedOffline(server1)
+		assert.Nil(t, err)
+		assert.True(t, isOffline)
+
+		err = db.RemoveServerAssumedOffline(server1)
+		assert.Nil(t, err)
+
+		isOffline, err = db.IsServerAssumedOffline(server1)
+		assert.Nil(t, err)
+		assert.False(t, isOffline)
+
+		err = db.SetServerAssumedOffline(server1)
+		assert.Nil(t, err)
+		err = db.SetServerAssumedOffline(server2)
+		assert.Nil(t, err)
+
+		isOffline, err = db.IsServerAssumedOffline(server1)
+		assert.Nil(t, err)
+		assert.True(t, isOffline)
+		isOffline, err = db.IsServerAssumedOffline(server2)
+		assert.Nil(t, err)
+		assert.True(t, isOffline)
+
+		err = db.RemoveAllServersAssumedOffline()
+		assert.Nil(t, err)
+
+		isOffline, err = db.IsServerAssumedOffline(server1)
+		assert.Nil(t, err)
+		assert.False(t, isOffline)
+		isOffline, err = db.IsServerAssumedOffline(server2)
+		assert.Nil(t, err)
+		assert.False(t, isOffline)
+	})
+}
+
+func TestRelayServersStored(t *testing.T) {
+	server := gomatrixserverlib.ServerName("server")
+	relayServer1 := gomatrixserverlib.ServerName("relayserver1")
+	relayServer2 := gomatrixserverlib.ServerName("relayserver2")
+
+	test.WithAllDatabases(t, func(t *testing.T, dbType test.DBType) {
+		db, closeDB := mustCreateFederationDatabase(t, dbType)
+		defer closeDB()
+
+		err := db.AddRelayServersForServer(server, []gomatrixserverlib.ServerName{relayServer1})
+		assert.Nil(t, err)
+
+		relayServers, err := db.GetRelayServersForServer(server)
+		assert.Nil(t, err)
+		assert.Equal(t, relayServer1, relayServers[0])
+
+		err = db.RemoveRelayServersForServer(server, []gomatrixserverlib.ServerName{relayServer1})
+		assert.Nil(t, err)
+
+		relayServers, err = db.GetRelayServersForServer(server)
+		assert.Nil(t, err)
+		assert.Zero(t, len(relayServers))
+
+		err = db.AddRelayServersForServer(server, []gomatrixserverlib.ServerName{relayServer1, relayServer2})
+		assert.Nil(t, err)
+
+		relayServers, err = db.GetRelayServersForServer(server)
+		assert.Nil(t, err)
+		assert.Equal(t, relayServer1, relayServers[0])
+		assert.Equal(t, relayServer2, relayServers[1])
+
+		err = db.RemoveAllRelayServersForServer(server)
+		assert.Nil(t, err)
+
+		relayServers, err = db.GetRelayServersForServer(server)
+		assert.Nil(t, err)
+		assert.Zero(t, len(relayServers))
+	})
+}