More safeguards when errors occur

matrix-org · Feb 12, 2024 · 212255a · 212255a
1 parent 66a31aa
commit 212255a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 13 deletions.
diff --git a/internal/api/rust/rust.go b/internal/api/rust/rust.go
@@ -214,20 +214,25 @@ func (c *RustClient) MustBackupKeys(t ct.TestLike) (recoveryKey string) {
 	genericListener := newGenericStateListener[matrix_sdk_ffi.EnableRecoveryProgress]()
 	var listener matrix_sdk_ffi.EnableRecoveryProgressListener = genericListener
 	recoveryKey, err := c.FFIClient.Encryption().EnableRecovery(true, listener)
-	for s := range genericListener.ch {
-		switch x := s.(type) {
-		case matrix_sdk_ffi.EnableRecoveryProgressCreatingBackup:
-			t.Logf("MustBackupKeys: state=CreatingBackup")
-		case matrix_sdk_ffi.EnableRecoveryProgressBackingUp:
-			t.Logf("MustBackupKeys: state=BackingUp %v/%v", x.BackedUpCount, x.TotalCount)
-		case matrix_sdk_ffi.EnableRecoveryProgressCreatingRecoveryKey:
-			t.Logf("MustBackupKeys: state=CreatingRecoveryKey")
-		case matrix_sdk_ffi.EnableRecoveryProgressDone:
-			t.Logf("MustBackupKeys: state=Done")
-			genericListener.Close() // break the loop
+	must.NotError(t, "Encryption.EnableRecovery", err)
+	for !genericListener.isClosed {
+		select {
+		case s := <-genericListener.ch:
+			switch x := s.(type) {
+			case matrix_sdk_ffi.EnableRecoveryProgressCreatingBackup:
+				t.Logf("MustBackupKeys: state=CreatingBackup")
+			case matrix_sdk_ffi.EnableRecoveryProgressBackingUp:
+				t.Logf("MustBackupKeys: state=BackingUp %v/%v", x.BackedUpCount, x.TotalCount)
+			case matrix_sdk_ffi.EnableRecoveryProgressCreatingRecoveryKey:
+				t.Logf("MustBackupKeys: state=CreatingRecoveryKey")
+			case matrix_sdk_ffi.EnableRecoveryProgressDone:
+				t.Logf("MustBackupKeys: state=Done")
+				genericListener.Close() // break the loop
+			}
+		case <-time.After(5 * time.Second):
+			ct.Fatalf(t, "timed out enabling backup keys")
 		}
 	}
-	must.NotError(t, "Encryption.EnableRecovery", err)
 	return recoveryKey
 }
 

diff --git a/tests/state_synchronisation_test.go b/tests/state_synchronisation_test.go
@@ -130,7 +130,11 @@ func testSigkillBeforeKeysUploadResponseJS(t *testing.T, clientType api.ClientTy
 		}
 		// destroy the client
 		mu.Lock()
-		terminateClient()
+		if terminateClient != nil {
+			terminateClient()
+		} else {
+			ct.Errorf(t, "terminateClient is nil. Did WithMITMOptions lock?")
+		}
 		mu.Unlock()
 	})
 	defer close()