Use MustUploadKeys more

Author: kegsay

Diff for: client/client.go

@@ -396,12 +396,18 @@ func (c *CSAPI) GetDefaultRoomVersion(t ct.TestLike) gomatrixserverlib.RoomVersi
 	return gomatrixserverlib.RoomVersion(defaultVersion.Str)
 }
 
+// MustUploadKeys uploads device and/or one time keys to the server, returning the current OTK counts.
+// Both device keys and one time keys are optional. Fails the test if the upload fails.
 func (c *CSAPI) MustUploadKeys(t ct.TestLike, deviceKeys map[string]interface{}, oneTimeKeys map[string]interface{}) (otkCounts map[string]int) {
 	t.Helper()
-	res := c.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"}, WithJSONBody(t, map[string]interface{}{
-		"device_keys":   deviceKeys,
-		"one_time_keys": oneTimeKeys,
-	}))
+	reqBody := make(map[string]interface{})
+	if deviceKeys != nil {
+		reqBody["device_keys"] = deviceKeys
+	}
+	if oneTimeKeys != nil {
+		reqBody["one_time_keys"] = oneTimeKeys
+	}
+	res := c.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"}, WithJSONBody(t, reqBody))
 	bodyBytes := ParseJSON(t, res)
 	s := struct {
 		OTKCounts map[string]int `json:"one_time_key_counts"`

Diff for: tests/csapi/keychanges_test.go

@@ -28,7 +28,8 @@ func TestKeyChangesLocal(t *testing.T) {
 	unauthedClient := deployment.UnauthenticatedClient(t, "hs1")
 
 	t.Run("New login should create a device_lists.changed entry", func(t *testing.T) {
-		mustUploadKeys(t, bob)
+		bobDeviceKeys, bobOTKs := bob.MustGenerateOneTimeKeys(t, 1)
+		bob.MustUploadKeys(t, bobDeviceKeys, bobOTKs)
 
 		roomID := alice.MustCreateRoom(t, map[string]interface{}{"preset": "public_chat"})
 		bob.MustJoinRoom(t, roomID, []string{})
@@ -48,7 +49,8 @@ func TestKeyChangesLocal(t *testing.T) {
 		unauthedClient.AccessToken = must.GetJSONFieldStr(t, loginResp, "access_token")
 		unauthedClient.DeviceID = must.GetJSONFieldStr(t, loginResp, "device_id")
 		unauthedClient.UserID = must.GetJSONFieldStr(t, loginResp, "user_id")
-		mustUploadKeys(t, unauthedClient)
+		unauthedKeys, unauthedOTKs := unauthedClient.MustGenerateOneTimeKeys(t, 1)
+		unauthedClient.MustUploadKeys(t, unauthedKeys, unauthedOTKs)
 
 		// Alice should now see a device list changed entry for Bob
 		nextBatch := alice.MustSyncUntil(t, client.SyncReq{Since: nextBatch1}, func(userID string, syncResp gjson.Result) error {
@@ -98,13 +100,3 @@ func TestKeyChangesLocal(t *testing.T) {
 		}
 	})
 }
-
-func mustUploadKeys(t *testing.T, user *client.CSAPI) {
-	t.Helper()
-	deviceKeys, oneTimeKeys := user.MustGenerateOneTimeKeys(t, 5)
-	reqBody := client.WithJSONBody(t, map[string]interface{}{
-		"device_keys":   deviceKeys,
-		"one_time_keys": oneTimeKeys,
-	})
-	user.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"}, reqBody)
-}

Diff for: tests/csapi/upload_keys_test.go

@@ -184,15 +184,12 @@ func TestKeyClaimOrdering(t *testing.T) {
 
 	// first upload key 1, sleep a bit, then upload key 0.
 	otk1 := map[string]interface{}{"signed_curve25519:1": oneTimeKeys["signed_curve25519:1"]}
-	alice.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"},
-		client.WithJSONBody(t, map[string]interface{}{"one_time_keys": otk1}))
-
-    // Ensure that there is a difference in timestamp between the two upload requests.
+	alice.MustUploadKeys(t, nil, otk1)
+	// Ensure that there is a difference in timestamp between the two upload requests.
 	time.Sleep(1 * time.Second)
 
 	otk0 := map[string]interface{}{"signed_curve25519:0": oneTimeKeys["signed_curve25519:0"]}
-	alice.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"},
-		client.WithJSONBody(t, map[string]interface{}{"one_time_keys": otk0}))
+	alice.MustUploadKeys(t, nil, otk0)
 
 	// Now claim the keys, and check they come back in the right order
 	reqBody := client.WithJSONBody(t, map[string]interface{}{

Diff for: tests/federation_device_list_update_test.go

@@ -142,10 +142,7 @@ func TestDeviceListsUpdateOverFederation(t *testing.T) {
 				Password: "this is alices password",
 			})
 			deviceKeys, oneTimeKeys := alice2.MustGenerateOneTimeKeys(t, 1)
-			alice2.MustDo(t, "POST", []string{"_matrix", "client", "v3", "keys", "upload"}, client.WithJSONBody(t, map[string]interface{}{
-				"device_keys":   deviceKeys,
-				"one_time_keys": oneTimeKeys,
-			}))
+			alice2.MustUploadKeys(t, deviceKeys, oneTimeKeys)
 
 			// now federation comes back online
 			tc.makeReachable(t)