From 19c669750ced0aec96277c4041a73c1742ea7a0f Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 5 Sep 2023 10:35:46 +0100
Subject: [PATCH] Fix flakey key upload test

---
 tests/federation_upload_keys_test.go | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/tests/federation_upload_keys_test.go b/tests/federation_upload_keys_test.go
index 4202bdbe..8baea823 100644
--- a/tests/federation_upload_keys_test.go
+++ b/tests/federation_upload_keys_test.go
@@ -22,6 +22,9 @@ func TestFederationKeyUploadQuery(t *testing.T) {
 	alice := deployment.Client(t, "hs1", "@alice:hs1")
 	bob := deployment.Client(t, "hs2", "@bob:hs2")
 
+	// Do an initial sync so that we can see the changes come down sync.
+	_, nextBatchBeforeKeyUpload := bob.MustSync(t, client.SyncReq{})
+
 	deviceKeys, oneTimeKeys := generateKeys(t, alice, 1)
 	// Upload keys
 	reqBody := client.WithJSONBody(t, map[string]interface{}{
@@ -80,9 +83,20 @@ func TestFederationKeyUploadQuery(t *testing.T) {
 
 	// sytest: Can query remote device keys using POST
 	t.Run("Can query remote device keys using POST", func(t *testing.T) {
-		// Device list changes only come down incremental syncs, so we do an
-		// initial sync up front.
-		_, nextBatch := bob.MustSync(t, client.SyncReq{})
+		// We expect the key upload to come down /sync. We need to do this so
+		// that can tell the next device update actually triggers the
+		// notification to go down /sync.
+		nextBatch := bob.MustSyncUntil(t, client.SyncReq{Since: nextBatchBeforeKeyUpload}, func(clientUserID string, topLevelSyncJSON gjson.Result) error {
+			devicesChanged := topLevelSyncJSON.Get("device_lists.changed")
+			if devicesChanged.Exists() {
+				for _, userID := range devicesChanged.Array() {
+					if userID.Str == alice.UserID {
+						return nil
+					}
+				}
+			}
+			return fmt.Errorf("no device_lists found")
+		})
 
 		displayName := "My new displayname"
 		body := client.WithJSONBody(t, map[string]interface{}{