Fix data race in unit tests

1 files changed, 10 insertions, 4 deletions

diff --git a/keyserver/internal/device_list_update_test.go b/keyserver/internal/device_list_update_test.go
index e8d1bfe8..59a66ec8 100644
--- a/keyserver/internal/device_list_update_test.go
+++ b/keyserver/internal/device_list_update_test.go
@@ -87,6 +87,12 @@ func (d *mockDeviceListUpdaterDatabase) MarkDeviceListStale(ctx context.Context,
 	return nil
 }
 
+func (d *mockDeviceListUpdaterDatabase) isStale(userID string) bool {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	return d.staleUsers[userID]
+}
+
 // StoreRemoteDeviceKeys persists the given keys. Keys with the same user ID and device ID will be replaced. An empty KeyJSON removes the key
 // for this (user, device). Does not modify the stream ID for keys.
 func (d *mockDeviceListUpdaterDatabase) StoreRemoteDeviceKeys(ctx context.Context, keys []api.DeviceMessage, clear []string) error {
@@ -169,7 +175,7 @@ func TestUpdateHavePrevID(t *testing.T) {
 	if !reflect.DeepEqual(db.storedKeys, []api.DeviceMessage{want}) {
 		t.Errorf("DB didn't store correct event, got %v want %v", db.storedKeys, want)
 	}
-	if db.staleUsers[event.UserID] {
+	if db.isStale(event.UserID) {
 		t.Errorf("%s incorrectly marked as stale", event.UserID)
 	}
 }
@@ -243,7 +249,7 @@ func TestUpdateNoPrevID(t *testing.T) {
 		},
 	}
 	// Now we should have a fresh list and the keys and emitted something
-	if db.staleUsers[event.UserID] {
+	if db.isStale(event.UserID) {
 		t.Errorf("%s still marked as stale", event.UserID)
 	}
 	if !reflect.DeepEqual(producer.events, []api.DeviceMessage{want}) {
@@ -304,7 +310,7 @@ func TestDebounce(t *testing.T) {
 	}
 
 	// user should be marked as stale
-	if !db.staleUsers[userID] {
+	if !db.isStale(userID) {
 		t.Errorf("user %s not marked as stale", userID)
 	}
 	// now send the response over federation
@@ -330,7 +336,7 @@ func TestDebounce(t *testing.T) {
 	wg.Wait()
 
 	// user is no longer stale now
-	if db.staleUsers[userID] {
+	if db.isStale(userID) {
 		t.Errorf("user %s is marked as stale", userID)
 	}
 }