Convert all tests to the new format

matrix-org · Jul 12, 2024 · 08044b2 · 08044b2
1 parent 5759abf
commit 08044b2
Show file tree

Hide file tree

Showing 5 changed files with 184 additions and 138 deletions.
diff --git a/internal/deploy/callback/callback_addon.go b/internal/deploy/callback/callback_addon.go
@@ -162,12 +162,13 @@ func NewCallbackServer(t ct.TestLike, hostnameRunningComplement string) (*Callba
 
 // SendError returns a callback.Fn which returns the provided statusCode
 // along with a JSON error $count times, after which it lets the response
-// pass through. This is useful for testing retries.
+// pass through. This is useful for testing retries. If count=0, always send
+// an error response.
 func SendError(count uint32, statusCode int) Fn {
 	var seen atomic.Uint32
 	return func(d Data) *Response {
 		next := seen.Add(1)
-		if next > count {
+		if count > 0 && next > count {
 			return nil
 		}
 		return &Response{

diff --git a/tests/notification_test.go b/tests/notification_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/matrix-org/complement-crypto/internal/api"
 	"github.com/matrix-org/complement-crypto/internal/cc"
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
+	"github.com/matrix-org/complement-crypto/internal/deploy/mitm"
 	"github.com/matrix-org/complement/must"
 )
 
@@ -567,23 +568,26 @@ func TestMultiprocessDupeOTKUpload(t *testing.T) {
 	// artificially slow down the HTTP responses, such that we will potentially have 2 in-flight /keys/upload requests
 	// at once. If the NSE and main apps are talking to each other, they should be using the same key ID + key.
 	// If not... well, that's a bug because then the client will forget one of these keys.
-	mitmConfiguration := tc.Deployment.MITM().Configure(t)
-	mitmConfiguration.ForPath("/keys/upload").Listen(func(cd callback.Data) *callback.Response {
-		if cd.AccessToken != aliceAccessToken {
-			return nil // let bob upload OTKs
-		}
-		aliceUploadedNewKeys = true
-		if cd.ResponseCode != 200 {
-			// we rely on the homeserver checking and rejecting when the same key ID is used with
-			// different keys.
-			t.Errorf("/keys/upload returned an error, duplicate key upload? %+v => %v", cd, string(cd.ResponseBody))
-		}
-		// tarpit the response
-		t.Logf("tarpitting keys/upload response for 4 seconds")
-		time.Sleep(4 * time.Second)
-		return nil
-	})
-	mitmConfiguration.Execute(func() {
+	tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+		Filter: mitm.FilterParams{
+			PathContains: "/keys/upload",
+		},
+		ResponseCallback: func(cd callback.Data) *callback.Response {
+			if cd.AccessToken != aliceAccessToken {
+				return nil // let bob upload OTKs
+			}
+			aliceUploadedNewKeys = true
+			if cd.ResponseCode != 200 {
+				// we rely on the homeserver checking and rejecting when the same key ID is used with
+				// different keys.
+				t.Errorf("/keys/upload returned an error, duplicate key upload? %+v => %v", cd, string(cd.ResponseBody))
+			}
+			// tarpit the response
+			t.Logf("tarpitting keys/upload response for 4 seconds")
+			time.Sleep(4 * time.Second)
+			return nil
+		},
+	}, func() {
 		var eventID string
 		// Bob appears and sends a message, causing Bob to claim one of Alice's OTKs.
 		// The main app will see this in /sync and then try to upload another OTK, which we will tarpit.

diff --git a/tests/one_time_keys_test.go b/tests/one_time_keys_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/matrix-org/complement-crypto/internal/api"
 	"github.com/matrix-org/complement-crypto/internal/cc"
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
+	"github.com/matrix-org/complement-crypto/internal/deploy/mitm"
 	"github.com/matrix-org/complement/b"
 	"github.com/matrix-org/complement/client"
 	"github.com/matrix-org/complement/ct"
@@ -108,9 +109,14 @@ func TestFallbackKeyIsUsedIfOneTimeKeysRunOut(t *testing.T) {
 			var roomID string
 			var waiter api.Waiter
 			// Block all /keys/upload requests for Alice
-			mitmConfiguration := tc.Deployment.MITM().Configure(t)
-			mitmConfiguration.ForPath("/keys/upload").AccessToken(alice.CurrentAccessToken(t)).BlockRequest(0, http.StatusGatewayTimeout)
-			mitmConfiguration.Execute(func() {
+			tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+				Filter: mitm.FilterParams{
+					PathContains: "/keys/upload",
+					Method:       "POST",
+					AccessToken:  alice.CurrentAccessToken(t),
+				},
+				RequestCallback: callback.SendError(0, http.StatusGatewayTimeout),
+			}, func() {
 				// claim all OTKs
 				mustClaimOTKs(t, otkGobbler, tc.Alice, int(otkCount))
 
@@ -156,9 +162,13 @@ func TestFailedOneTimeKeyUploadRetries(t *testing.T) {
 		// make a room so we can kick clients
 		roomID := tc.Alice.MustCreateRoom(t, map[string]interface{}{"preset": "public_chat"})
 		// block /keys/upload and make a client
-		mitmConfiguration := tc.Deployment.MITM().Configure(t)
-		mitmConfiguration.ForPath("/keys/upload").Method("POST").BlockRequest(2, http.StatusGatewayTimeout)
-		mitmConfiguration.Execute(func() {
+		tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+			Filter: mitm.FilterParams{
+				PathContains: "/keys/upload",
+				Method:       "POST",
+			},
+			RequestCallback: callback.SendError(2, http.StatusGatewayTimeout),
+		}, func() {
 			tc.WithAliceSyncing(t, func(alice api.Client) {
 				tc.Bob.MustDo(t, "POST", []string{
 					"_matrix", "client", "v3", "keys", "claim",
@@ -204,16 +214,21 @@ func TestFailedKeysClaimRetries(t *testing.T) {
 			// make a room which will link the 2 users together when
 			roomID := tc.CreateNewEncryptedRoom(t, tc.Alice, cc.EncRoomOptions.PresetPublicChat())
 			// block /keys/claim and join the room, causing the Olm session to be created
-			mitmConfiguration := tc.Deployment.MITM().Configure(t)
-			mitmConfiguration.ForPath("/keys/claim").Method("POST").BlockRequest(2, http.StatusGatewayTimeout).Listen(func(cd callback.Data) *callback.Response {
-				t.Logf("%+v", cd)
-				if cd.ResponseCode == 200 {
-					waiter.Finish()
-					stopPoking.Store(true)
-				}
-				return nil
-			})
-			mitmConfiguration.Execute(func() {
+			tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+				Filter: mitm.FilterParams{
+					PathContains: "/keys/claim",
+					Method:       "POST",
+				},
+				RequestCallback: callback.SendError(2, http.StatusGatewayTimeout),
+				ResponseCallback: func(cd callback.Data) *callback.Response {
+					t.Logf("%+v", cd)
+					if cd.ResponseCode == 200 {
+						waiter.Finish()
+						stopPoking.Store(true)
+					}
+					return nil
+				},
+			}, func() {
 				// join the room. This should cause an Olm session to be made but it will fail as we cannot
 				// call /keys/claim. We should retry though.
 				tc.Bob.MustJoinRoom(t, roomID, []string{clientType.HS})

diff --git a/tests/state_synchronisation_test.go b/tests/state_synchronisation_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/matrix-org/complement-crypto/internal/api"
 	"github.com/matrix-org/complement-crypto/internal/cc"
 	"github.com/matrix-org/complement-crypto/internal/deploy/callback"
+	"github.com/matrix-org/complement-crypto/internal/deploy/mitm"
 	"github.com/matrix-org/complement/ct"
 	"github.com/matrix-org/complement/helpers"
 )
@@ -36,25 +37,28 @@ func testSigkillBeforeKeysUploadResponseRust(t *testing.T, clientType api.Client
 	var terminateClient func()
 	seenSecondKeysUploadWaiter := helpers.NewWaiter()
 	tc := Instance().CreateTestContext(t, clientType, clientType)
-
-	mitmConfiguration := tc.Deployment.MITM().Configure(t)
-	mitmConfiguration.ForPath("/keys/upload").Listen(func(cd callback.Data) *callback.Response {
-		if terminated.Load() {
-			// make sure the 2nd upload 200 OKs
-			if cd.ResponseCode != 200 {
-				t.Errorf("2nd /keys/upload did not 200 OK => got %v", cd.ResponseCode)
+	tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+		Filter: mitm.FilterParams{
+			PathContains: "/keys/upload",
+			Method:       "POST",
+		},
+		ResponseCallback: func(cd callback.Data) *callback.Response {
+			if terminated.Load() {
+				// make sure the 2nd upload 200 OKs
+				if cd.ResponseCode != 200 {
+					t.Errorf("2nd /keys/upload did not 200 OK => got %v", cd.ResponseCode)
+				}
+				t.Logf("recv 2nd /keys/upload => HTTP %d", cd.ResponseCode)
+				seenSecondKeysUploadWaiter.Finish()
+				return nil
 			}
-			t.Logf("recv 2nd /keys/upload => HTTP %d", cd.ResponseCode)
-			seenSecondKeysUploadWaiter.Finish()
+			// destroy the client
+			mu.Lock()
+			terminateClient()
+			mu.Unlock()
 			return nil
-		}
-		// destroy the client
-		mu.Lock()
-		terminateClient()
-		mu.Unlock()
-		return nil
-	})
-	mitmConfiguration.Execute(func() {
+		},
+	}, func() {
 		// login in a different process
 		remoteClient := tc.MustCreateClient(t, &cc.ClientCreationRequest{
 			User: tc.Alice,
@@ -98,30 +102,31 @@ func testSigkillBeforeKeysUploadResponseJS(t *testing.T, clientType api.ClientTy
 	var terminateClient func()
 	seenSecondKeysUploadWaiter := helpers.NewWaiter()
 	tc := Instance().CreateTestContext(t, clientType, clientType)
-	mitmConfiguration := tc.Deployment.MITM().Configure(t)
-	mitmConfiguration.ForPath("/keys/upload").Listen(func(cd callback.Data) *callback.Response {
-		if cd.Method == "OPTIONS" {
-			return nil // ignore CORS
-		}
-		if terminated.Load() {
-			// make sure the 2nd upload 200 OKs
-			if cd.ResponseCode != 200 {
-				ct.Errorf(t, "2nd /keys/upload did not 200 OK => got %v", cd.ResponseCode)
+	tc.Deployment.MITM().Configure(t).WithIntercept(mitm.InterceptOpts{
+		Filter: mitm.FilterParams{
+			PathContains: "/keys/upload",
+			Method:       "POST",
+		},
+		ResponseCallback: func(cd callback.Data) *callback.Response {
+			if terminated.Load() {
+				// make sure the 2nd upload 200 OKs
+				if cd.ResponseCode != 200 {
+					ct.Errorf(t, "2nd /keys/upload did not 200 OK => got %v", cd.ResponseCode)
+				}
+				seenSecondKeysUploadWaiter.Finish()
+				return nil
 			}
-			seenSecondKeysUploadWaiter.Finish()
+			// destroy the client
+			mu.Lock()
+			if terminateClient != nil {
+				terminateClient()
+			} else {
+				ct.Errorf(t, "terminateClient is nil. Did WithMITMOptions lock?")
+			}
+			mu.Unlock()
 			return nil
-		}
-		// destroy the client
-		mu.Lock()
-		if terminateClient != nil {
-			terminateClient()
-		} else {
-			ct.Errorf(t, "terminateClient is nil. Did WithMITMOptions lock?")
-		}
-		mu.Unlock()
-		return nil
-	})
-	mitmConfiguration.Execute(func() {
+		},
+	}, func() {
 		clientWhichWillBeKilled := tc.MustCreateClient(t, &cc.ClientCreationRequest{
 			User: tc.Alice,
 			Opts: api.ClientCreationOpts{