Add Must[Get|Set]DisplayName and use MustUploadKeys in more places (#755)

* Add MustGetDisplayName and MustSetDisplayName

* Use MustUploadKeys more

Author: kegsay

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"`
@@ -492,6 +498,20 @@ func (c *CSAPI) MustGenerateOneTimeKeys(t ct.TestLike, otkCount uint) (deviceKey
 	return deviceKeys, oneTimeKeys
 }
 
+// MustSetDisplayName sets the global display name for this account or fails the test.
+func (c *CSAPI) MustSetDisplayName(t ct.TestLike, displayname string) {
+	c.MustDo(t, "PUT", []string{"_matrix", "client", "v3", "profile", c.UserID, "displayname"}, WithJSONBody(t, map[string]any{
+		"displayname": displayname,
+	}))
+}
+
+// MustGetDisplayName returns the global display name for this user or fails the test.
+func (c *CSAPI) MustGetDisplayName(t ct.TestLike, userID string) string {
+	res := c.MustDo(t, "GET", []string{"_matrix", "client", "v3", "profile", userID, "displayname"})
+	body := ParseJSON(t, res)
+	return GetJSONFieldStr(t, body, "displayname")
+}
+
 // WithRawBody sets the HTTP request body to `body`
 func WithRawBody(body []byte) RequestOpt {
 	return func(req *http.Request) {

tests/csapi/apidoc_profile_displayname_test.go

@@ -4,9 +4,7 @@ import (
 	"testing"
 
 	"github.com/matrix-org/complement"
-	"github.com/matrix-org/complement/client"
 	"github.com/matrix-org/complement/helpers"
-	"github.com/matrix-org/complement/match"
 	"github.com/matrix-org/complement/must"
 )
 
@@ -18,19 +16,11 @@ func TestProfileDisplayName(t *testing.T) {
 	displayName := "my_display_name"
 	// sytest: PUT /profile/:user_id/displayname sets my name
 	t.Run("PUT /profile/:user_id/displayname sets my name", func(t *testing.T) {
-		reqBody := client.WithJSONBody(t, map[string]interface{}{
-			"displayname": displayName,
-		})
-		_ = authedClient.MustDo(t, "PUT", []string{"_matrix", "client", "v3", "profile", authedClient.UserID, "displayname"}, reqBody)
+		authedClient.MustSetDisplayName(t, displayName)
 	})
 	// sytest: GET /profile/:user_id/displayname publicly accessible
 	t.Run("GET /profile/:user_id/displayname publicly accessible", func(t *testing.T) {
-		res := unauthedClient.Do(t, "GET", []string{"_matrix", "client", "v3", "profile", authedClient.UserID, "displayname"})
-		must.MatchResponse(t, res, match.HTTPResponse{
-			StatusCode: 200,
-			JSON: []match.JSON{
-				match.JSONKeyEqual("displayname", displayName),
-			},
-		})
+		gotDisplayName := unauthedClient.MustGetDisplayName(t, authedClient.UserID)
+		must.Equal(t, gotDisplayName, displayName, "display name mismatch")
 	})
 }

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)
-}

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{}{

tests/csapi/user_directory_display_names_test.go

@@ -61,14 +61,7 @@ func setupUsers(t *testing.T) (*client.CSAPI, *client.CSAPI, *client.CSAPI, func
 	// Alice sets her profile displayname. This ensures that her
 	// public name, private name and userid localpart are all
 	// distinguishable, even case-insensitively.
-	alice.MustDo(
-		t,
-		"PUT",
-		[]string{"_matrix", "client", "v3", "profile", alice.UserID, "displayname"},
-		client.WithJSONBody(t, map[string]interface{}{
-			"displayname": alicePublicName,
-		}),
-	)
+	alice.MustSetDisplayName(t, alicePublicName)
 
 	// Alice creates a public room (so when Eve searches, she can see that Alice exists)
 	alice.MustCreateRoom(t, map[string]interface{}{"visibility": "public"})

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)

tests/federation_query_profile_test.go

@@ -10,9 +10,8 @@ import (
 	"github.com/matrix-org/gomatrixserverlib/fclient"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 
-	"github.com/matrix-org/complement/client"
-	"github.com/matrix-org/complement/helpers"
 	"github.com/matrix-org/complement/federation"
+	"github.com/matrix-org/complement/helpers"
 	"github.com/matrix-org/complement/match"
 	"github.com/matrix-org/complement/must"
 )
@@ -61,12 +60,8 @@ func TestOutboundFederationProfile(t *testing.T) {
 
 		// query the display name which should do an outbound federation hit
 		unauthedClient := deployment.UnauthenticatedClient(t, "hs1")
-		res := unauthedClient.MustDo(t, "GET", []string{"_matrix", "client", "v3", "profile", remoteUserID, "displayname"})
-		must.MatchResponse(t, res, match.HTTPResponse{
-			JSON: []match.JSON{
-				match.JSONKeyEqual("displayname", remoteDisplayName),
-			},
-		})
+		gotDisplayName := unauthedClient.MustGetDisplayName(t, remoteUserID)
+		must.Equal(t, gotDisplayName, remoteDisplayName, "display name mismatch")
 	})
 }
 
@@ -107,14 +102,7 @@ func TestInboundFederationProfile(t *testing.T) {
 	t.Run("Inbound federation can query profile data", func(t *testing.T) {
 		const alicePublicName = "Alice Cooper"
 
-		alice.MustDo(
-			t,
-			"PUT",
-			[]string{"_matrix", "client", "v3", "profile", alice.UserID, "displayname"},
-			client.WithJSONBody(t, map[string]interface{}{
-				"displayname": alicePublicName,
-			}),
-		)
+		alice.MustSetDisplayName(t, alicePublicName)
 
 		fedReq := fclient.NewFederationRequest(
 			"GET",

tests/msc3902/federation_room_join_partial_state_test.go

@@ -3597,13 +3597,7 @@ func TestPartialStateJoin(t *testing.T) {
 		)
 		defer removePDUHandler()
 
-		alice.MustDo(t,
-			"PUT",
-			[]string{"_matrix", "client", "v3", "profile", alice.UserID, "displayname"},
-			client.WithJSONBody(t, map[string]interface{}{
-				"displayname": "alice 2",
-			}),
-		)
+		alice.MustSetDisplayName(t, "alice 2")
 		t.Logf("Alice changed display name")
 
 		select {