Add support for PNI registration ids and PNP change number.

app/src/androidTest/java/org/thoughtcrime/securesms/SignalInstrumentationApplicationContext.kt

@@ -0,0 +1,15 @@
+package org.thoughtcrime.securesms
+
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencyProvider
+import org.thoughtcrime.securesms.dependencies.InstrumentationApplicationDependencyProvider
+
+/**
+ * Application context for running instrumentation tests (aka androidTests).
+ */
+class SignalInstrumentationApplicationContext : ApplicationContext() {
+  override fun initializeAppDependencies() {
+    val default = ApplicationDependencyProvider(this)
+    ApplicationDependencies.init(this, InstrumentationApplicationDependencyProvider(this, default))
+  }
+}

app/src/androidTest/java/org/thoughtcrime/securesms/components/settings/app/changenumber/ChangeNumberViewModelTest.kt

@@ -0,0 +1,272 @@
+package org.thoughtcrime.securesms.components.settings.app.changenumber
+
+import androidx.lifecycle.SavedStateHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import okhttp3.mockwebserver.MockResponse
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.signal.core.util.ThreadUtil
+import org.signal.libsignal.protocol.state.SignedPreKeyRecord
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies
+import org.thoughtcrime.securesms.dependencies.InstrumentationApplicationDependencyProvider
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.pin.KbsRepository
+import org.thoughtcrime.securesms.recipients.Recipient
+import org.thoughtcrime.securesms.registration.VerifyAccountRepository
+import org.thoughtcrime.securesms.registration.VerifyAccountResponseProcessor
+import org.thoughtcrime.securesms.testing.Get
+import org.thoughtcrime.securesms.testing.MockProvider
+import org.thoughtcrime.securesms.testing.Put
+import org.thoughtcrime.securesms.testing.SignalActivityRule
+import org.thoughtcrime.securesms.testing.assertIs
+import org.thoughtcrime.securesms.testing.assertIsNot
+import org.thoughtcrime.securesms.testing.assertIsNotNull
+import org.thoughtcrime.securesms.testing.assertIsNull
+import org.thoughtcrime.securesms.testing.assertIsSize
+import org.thoughtcrime.securesms.testing.connectionFailure
+import org.thoughtcrime.securesms.testing.failure
+import org.thoughtcrime.securesms.testing.parsedRequestBody
+import org.thoughtcrime.securesms.testing.success
+import org.thoughtcrime.securesms.testing.timeout
+import org.whispersystems.signalservice.api.account.ChangePhoneNumberRequest
+import org.whispersystems.signalservice.api.push.ServiceId
+import org.whispersystems.signalservice.internal.push.PreKeyState
+import java.util.UUID
+
+@RunWith(AndroidJUnit4::class)
+class ChangeNumberViewModelTest {
+
+  @get:Rule
+  val harness = SignalActivityRule()
+
+  private lateinit var viewModel: ChangeNumberViewModel
+  private lateinit var kbsRepository: KbsRepository
+
+  @Before
+  fun setUp() {
+    ApplicationDependencies.getSignalServiceAccountManager().setSoTimeoutMillis(1000)
+    kbsRepository = mock()
+    ThreadUtil.runOnMainSync {
+      viewModel = ChangeNumberViewModel(
+        localNumber = harness.self.requireE164(),
+        changeNumberRepository = ChangeNumberRepository(),
+        savedState = SavedStateHandle(),
+        password = SignalStore.account().servicePassword!!,
+        verifyAccountRepository = VerifyAccountRepository(harness.application),
+        kbsRepository = kbsRepository
+      )
+
+      viewModel.setNewCountry(1)
+      viewModel.setNewNationalNumber("5555550102")
+    }
+  }
+
+  @After
+  fun tearDown() {
+    InstrumentationApplicationDependencyProvider.clearHandlers()
+  }
+
+  @Test
+  fun testChangeNumber_givenOnlyPrimaryAndNoRegLock() {
+    // GIVEN
+    val aci = Recipient.self().requireServiceId()
+    val newPni = ServiceId.from(UUID.randomUUID())
+    lateinit var changeNumberRequest: ChangePhoneNumberRequest
+    lateinit var setPreKeysRequest: PreKeyState
+
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(
+      Get("/v1/devices") { MockResponse().success(MockProvider.primaryOnlyDeviceList) },
+      Put("/v1/accounts/number") { r ->
+        changeNumberRequest = r.parsedRequestBody()
+        MockResponse().success(MockProvider.createVerifyAccountResponse(aci, newPni))
+      },
+      Put("/v2/keys") { r ->
+        setPreKeysRequest = r.parsedRequestBody()
+        MockResponse().success()
+      },
+      Get("/v1/certificate/delivery") { MockResponse().success(MockProvider.senderCertificate) }
+    )
+
+    // WHEN
+    viewModel.verifyCodeWithoutRegistrationLock("123456").blockingGet().resultOrThrow
+
+    // THEN
+    assertSuccess(newPni, changeNumberRequest, setPreKeysRequest)
+  }
+
+  /**
+   * If we encounter a server error, this means the server ack our request and rejected it. In this
+   * case we know the change *did not* take on the server and can reset to a clean state.
+   */
+  @Test
+  fun testChangeNumber_givenServerFailedApiCall() {
+    // GIVEN
+    val oldPni = Recipient.self().requirePni()
+    val oldE164 = Recipient.self().requireE164()
+
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(
+      Get("/v1/devices") { MockResponse().success(MockProvider.primaryOnlyDeviceList) },
+      Put("/v1/accounts/number") { MockResponse().failure(500) },
+    )
+
+    // WHEN
+    val processor: VerifyAccountResponseProcessor = viewModel.verifyCodeWithoutRegistrationLock("123456").blockingGet()
+
+    // THEN
+    processor.isServerSentError() assertIs true
+    Recipient.self().requireE164() assertIs oldE164
+    Recipient.self().requirePni() assertIs oldPni
+    SignalStore.misc().pendingChangeNumberMetadata.assertIsNull()
+  }
+
+  /**
+   * If we encounter a non-server error like a timeout or bad SSL, we do not know the state of our change
+   * number on the server side. We have to do a whoami call to query the server for our details and then
+   * respond accordingly.
+   *
+   * In this case, the whoami is our old details, so we can know the change *did not* take on the server
+   * and can reset to a clean state.
+   */
+  @Test
+  fun testChangeNumber_givenNetworkFailedApiCallEnRouteToServer() {
+    // GIVEN
+    val aci = Recipient.self().requireServiceId()
+    val oldPni = Recipient.self().requirePni()
+    val oldE164 = Recipient.self().requireE164()
+
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(
+      Get("/v1/devices") { MockResponse().success(MockProvider.primaryOnlyDeviceList) },
+      Put("/v1/accounts/number") { MockResponse().connectionFailure() },
+      Get("/v1/accounts/whoami") { MockResponse().success(MockProvider.createWhoAmIResponse(aci, oldPni, oldE164)) }
+    )
+
+    // WHEN
+    val processor: VerifyAccountResponseProcessor = viewModel.verifyCodeWithoutRegistrationLock("123456").blockingGet()
+
+    // THEN
+    processor.isServerSentError() assertIs false
+    Recipient.self().requireE164() assertIs oldE164
+    Recipient.self().requirePni() assertIs oldPni
+    SignalStore.misc().isChangeNumberLocked assertIs false
+    SignalStore.misc().pendingChangeNumberMetadata.assertIsNull()
+  }
+
+  /**
+   * If we encounter a non-server error like a timeout or bad SSL, we do not know the state of our change
+   * number on the server side. We have to do a whoami call to query the server for our details and then
+   * respond accordingly.
+   *
+   * In this case, the whoami is our new details, so we can know the change *did* take on the server
+   * and need to keep the app in a locked state. The test then uses the ChangeNumberLockActivity to unlock
+   * and apply the pending state after confirming the change on the server.
+   */
+  @Test
+  fun testChangeNumber_givenNetworkFailedApiCallEnRouteToClient() {
+    // GIVEN
+    val aci = Recipient.self().requireServiceId()
+    val oldPni = Recipient.self().requirePni()
+    val oldE164 = Recipient.self().requireE164()
+    val newPni = ServiceId.from(UUID.randomUUID())
+
+    lateinit var changeNumberRequest: ChangePhoneNumberRequest
+    lateinit var setPreKeysRequest: PreKeyState
+
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(
+      Get("/v1/devices") { MockResponse().success(MockProvider.primaryOnlyDeviceList) },
+      Put("/v1/accounts/number") { r ->
+        changeNumberRequest = r.parsedRequestBody()
+        MockResponse().timeout()
+      },
+      Get("/v1/accounts/whoami") { MockResponse().success(MockProvider.createWhoAmIResponse(aci, newPni, "+15555550102")) },
+      Put("/v2/keys") { r ->
+        setPreKeysRequest = r.parsedRequestBody()
+        MockResponse().success()
+      },
+      Get("/v1/certificate/delivery") { MockResponse().success(MockProvider.senderCertificate) }
+    )
+
+    // WHEN
+    val processor: VerifyAccountResponseProcessor = viewModel.verifyCodeWithoutRegistrationLock("123456").blockingGet()
+
+    // THEN
+    processor.isServerSentError() assertIs false
+    Recipient.self().requireE164() assertIs oldE164
+    Recipient.self().requirePni() assertIs oldPni
+    SignalStore.misc().isChangeNumberLocked assertIs true
+    SignalStore.misc().pendingChangeNumberMetadata.assertIsNotNull()
+
+    // WHEN AGAIN Processing lock
+    val scenario = harness.launchActivity<ChangeNumberLockActivity>()
+    scenario.onActivity {}
+    ThreadUtil.sleep(500)
+
+    // THEN AGAIN
+    assertSuccess(newPni, changeNumberRequest, setPreKeysRequest)
+  }
+
+  @Test
+  fun testChangeNumber_givenOnlyPrimaryAndRegistrationLock() {
+    // GIVEN
+    val aci = Recipient.self().requireServiceId()
+    val newPni = ServiceId.from(UUID.randomUUID())
+
+    lateinit var changeNumberRequest: ChangePhoneNumberRequest
+    lateinit var setPreKeysRequest: PreKeyState
+
+    MockProvider.mockGetRegistrationLockStringFlow(kbsRepository)
+
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(
+      Get("/v1/devices") { MockResponse().success(MockProvider.primaryOnlyDeviceList) },
+      Put("/v1/accounts/number") { r ->
+        changeNumberRequest = r.parsedRequestBody()
+        if (changeNumberRequest.registrationLock.isNullOrEmpty()) {
+          MockResponse().failure(423, MockProvider.lockedFailure)
+        } else {
+          MockResponse().success(MockProvider.createVerifyAccountResponse(aci, newPni))
+        }
+      },
+      Put("/v2/keys") { r ->
+        setPreKeysRequest = r.parsedRequestBody()
+        MockResponse().success()
+      },
+      Get("/v1/certificate/delivery") { MockResponse().success(MockProvider.senderCertificate) }
+    )
+
+    // WHEN
+    viewModel.verifyCodeWithoutRegistrationLock("123456").blockingGet().also { processor ->
+      processor.registrationLock() assertIs true
+      Recipient.self().requirePni() assertIsNot newPni
+      SignalStore.misc().pendingChangeNumberMetadata.assertIsNull()
+    }
+
+    viewModel.verifyCodeAndRegisterAccountWithRegistrationLock("pin").blockingGet().resultOrThrow
+
+    // THEN
+    assertSuccess(newPni, changeNumberRequest, setPreKeysRequest)
+  }
+
+  private fun assertSuccess(newPni: ServiceId, changeNumberRequest: ChangePhoneNumberRequest, setPreKeysRequest: PreKeyState) {
+    val pniProtocolStore = ApplicationDependencies.getProtocolStore().pni()
+    val pniMetadataStore = SignalStore.account().pniPreKeys
+
+    Recipient.self().requireE164() assertIs "+15555550102"
+    Recipient.self().requirePni() assertIs newPni
+
+    SignalStore.account().pniRegistrationId assertIs changeNumberRequest.pniRegistrationIds["1"]!!
+    SignalStore.account().pniIdentityKey.publicKey assertIs changeNumberRequest.pniIdentityKey
+    pniMetadataStore.activeSignedPreKeyId assertIs changeNumberRequest.devicePniSignedPrekeys["1"]!!.keyId
+
+    val activeSignedPreKey: SignedPreKeyRecord = pniProtocolStore.loadSignedPreKey(pniMetadataStore.activeSignedPreKeyId)
+    activeSignedPreKey.keyPair.publicKey assertIs changeNumberRequest.devicePniSignedPrekeys["1"]!!.publicKey
+    activeSignedPreKey.signature assertIs changeNumberRequest.devicePniSignedPrekeys["1"]!!.signature
+
+    setPreKeysRequest.signedPreKey.publicKey assertIs activeSignedPreKey.keyPair.publicKey
+    setPreKeysRequest.preKeys assertIsSize 100
+
+    SignalStore.misc().pendingChangeNumberMetadata.assertIsNull()
+  }
+}

app/src/androidTest/java/org/thoughtcrime/securesms/dependencies/InstrumentationApplicationDependencyProvider.kt

@@ -0,0 +1,109 @@
+package org.thoughtcrime.securesms.dependencies
+
+import android.app.Application
+import okhttp3.ConnectionSpec
+import okhttp3.mockwebserver.Dispatcher
+import okhttp3.mockwebserver.MockResponse
+import okhttp3.mockwebserver.MockWebServer
+import okhttp3.mockwebserver.RecordedRequest
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.thoughtcrime.securesms.BuildConfig
+import org.thoughtcrime.securesms.KbsEnclave
+import org.thoughtcrime.securesms.push.SignalServiceNetworkAccess
+import org.thoughtcrime.securesms.push.SignalServiceTrustStore
+import org.thoughtcrime.securesms.testing.Verb
+import org.thoughtcrime.securesms.testing.runSync
+import org.thoughtcrime.securesms.util.Base64
+import org.whispersystems.signalservice.api.KeyBackupService
+import org.whispersystems.signalservice.api.SignalServiceAccountManager
+import org.whispersystems.signalservice.api.push.TrustStore
+import org.whispersystems.signalservice.internal.configuration.SignalCdnUrl
+import org.whispersystems.signalservice.internal.configuration.SignalCdsiUrl
+import org.whispersystems.signalservice.internal.configuration.SignalContactDiscoveryUrl
+import org.whispersystems.signalservice.internal.configuration.SignalKeyBackupServiceUrl
+import org.whispersystems.signalservice.internal.configuration.SignalServiceConfiguration
+import org.whispersystems.signalservice.internal.configuration.SignalServiceUrl
+import org.whispersystems.signalservice.internal.configuration.SignalStorageUrl
+import java.security.KeyStore
+import java.util.Optional
+
+/**
+ * Dependency provider used for instrumentation tests (aka androidTests).
+ *
+ * Handles setting up a mock web server for API calls, and provides mockable versions of [SignalServiceNetworkAccess] and
+ * [KeyBackupService].
+ */
+class InstrumentationApplicationDependencyProvider(application: Application, default: ApplicationDependencyProvider) : ApplicationDependencies.Provider by default {
+
+  private val serviceTrustStore: TrustStore
+  private val uncensoredConfiguration: SignalServiceConfiguration
+  private val serviceNetworkAccessMock: SignalServiceNetworkAccess
+  private val keyBackupService: KeyBackupService
+
+  init {
+    runSync {
+      webServer = MockWebServer()
+      baseUrl = webServer.url("").toString()
+    }
+
+    webServer.setDispatcher(object : Dispatcher() {
+      override fun dispatch(request: RecordedRequest): MockResponse {
+        val handler = handlers.firstOrNull {
+          request.method == it.verb && request.path.startsWith("/${it.path}")
+        }
+        return handler?.responseFactory?.invoke(request) ?: MockResponse().setResponseCode(500)
+      }
+    })
+
+    serviceTrustStore = SignalServiceTrustStore(application)
+    uncensoredConfiguration = SignalServiceConfiguration(
+      arrayOf(SignalServiceUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+      mapOf(
+        0 to arrayOf(SignalCdnUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+        2 to arrayOf(SignalCdnUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT))
+      ),
+      arrayOf(SignalContactDiscoveryUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+      arrayOf(SignalKeyBackupServiceUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+      arrayOf(SignalStorageUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+      arrayOf(SignalCdsiUrl(baseUrl, "localhost", serviceTrustStore, ConnectionSpec.CLEARTEXT)),
+      emptyList(),
+      Optional.of(SignalServiceNetworkAccess.DNS),
+      Optional.empty(),
+      Base64.decode(BuildConfig.ZKGROUP_SERVER_PUBLIC_PARAMS)
+    )
+
+    serviceNetworkAccessMock = mock {
+      on { getConfiguration() } doReturn uncensoredConfiguration
+      on { getConfiguration(any()) } doReturn uncensoredConfiguration
+    }
+
+    keyBackupService = mock()
+  }
+
+  override fun provideSignalServiceNetworkAccess(): SignalServiceNetworkAccess {
+    return serviceNetworkAccessMock
+  }
+
+  override fun provideKeyBackupService(signalServiceAccountManager: SignalServiceAccountManager, keyStore: KeyStore, enclave: KbsEnclave): KeyBackupService {
+    return keyBackupService
+  }
+
+  companion object {
+    lateinit var webServer: MockWebServer
+      private set
+    lateinit var baseUrl: String
+      private set
+
+    private val handlers: MutableList<Verb> = mutableListOf()
+
+    fun addMockWebRequestHandlers(vararg verbs: Verb) {
+      handlers.addAll(verbs)
+    }
+
+    fun clearHandlers() {
+      handlers.clear()
+    }
+  }
+}

app/src/androidTest/java/org/thoughtcrime/securesms/testing/MockProvider.kt

@@ -0,0 +1,86 @@
+package org.thoughtcrime.securesms.testing
+
+import io.reactivex.rxjava3.core.Single
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.stub
+import org.signal.core.util.Hex
+import org.thoughtcrime.securesms.dependencies.ApplicationDependencies
+import org.thoughtcrime.securesms.pin.KbsRepository
+import org.thoughtcrime.securesms.pin.TokenData
+import org.thoughtcrime.securesms.test.BuildConfig
+import org.whispersystems.signalservice.api.KbsPinData
+import org.whispersystems.signalservice.api.KeyBackupService
+import org.whispersystems.signalservice.api.kbs.HashedPin
+import org.whispersystems.signalservice.api.kbs.MasterKey
+import org.whispersystems.signalservice.api.messages.multidevice.DeviceInfo
+import org.whispersystems.signalservice.api.push.ServiceId
+import org.whispersystems.signalservice.internal.ServiceResponse
+import org.whispersystems.signalservice.internal.contacts.entities.TokenResponse
+import org.whispersystems.signalservice.internal.push.AuthCredentials
+import org.whispersystems.signalservice.internal.push.DeviceInfoList
+import org.whispersystems.signalservice.internal.push.PushServiceSocket
+import org.whispersystems.signalservice.internal.push.SenderCertificate
+import org.whispersystems.signalservice.internal.push.VerifyAccountResponse
+import org.whispersystems.signalservice.internal.push.WhoAmIResponse
+import java.security.SecureRandom
+
+/**
+ * Warehouse of reusable test data and mock configurations.
+ */
+object MockProvider {
+
+  val senderCertificate = SenderCertificate().apply { certificate = ByteArray(0) }
+
+  val lockedFailure = PushServiceSocket.RegistrationLockFailure().apply {
+    backupCredentials = AuthCredentials.create("username", "password")
+  }
+
+  val primaryOnlyDeviceList = DeviceInfoList().apply {
+    devices = listOf(
+      DeviceInfo().apply {
+        id = 1
+      }
+    )
+  }
+
+  fun createVerifyAccountResponse(aci: ServiceId, newPni: ServiceId): VerifyAccountResponse {
+    return VerifyAccountResponse().apply {
+      uuid = aci.toString()
+      pni = newPni.toString()
+      storageCapable = false
+    }
+  }
+
+  fun createWhoAmIResponse(aci: ServiceId, pni: ServiceId, e164: String): WhoAmIResponse {
+    return WhoAmIResponse().apply {
+      this.uuid = aci.toString()
+      this.pni = pni.toString()
+      this.number = e164
+    }
+  }
+
+  fun mockGetRegistrationLockStringFlow(kbsRepository: KbsRepository) {
+    val tokenData: TokenData = mock {
+      on { enclave } doReturn BuildConfig.KBS_ENCLAVE
+      on { basicAuth } doReturn "basicAuth"
+      on { triesRemaining } doReturn 10
+      on { tokenResponse } doReturn TokenResponse()
+    }
+
+    kbsRepository.stub {
+      on { getToken(any() as String) } doReturn Single.just(ServiceResponse.forResult(tokenData, 200, ""))
+    }
+
+    val session: KeyBackupService.RestoreSession = object : KeyBackupService.RestoreSession {
+      override fun hashSalt(): ByteArray = Hex.fromStringCondensed("cba811749042b303a6a7efa5ccd160aea5e3ea243c8d2692bd13d515732f51a8")
+      override fun restorePin(hashedPin: HashedPin?): KbsPinData = KbsPinData(MasterKey.createNew(SecureRandom()), null)
+    }
+
+    val kbsService = ApplicationDependencies.getKeyBackupService(BuildConfig.KBS_ENCLAVE)
+    kbsService.stub {
+      on { newRegistrationSession(any(), any()) } doReturn session
+    }
+  }
+}

app/src/androidTest/java/org/thoughtcrime/securesms/testing/ResponseMocking.kt

@@ -0,0 +1,48 @@
+package org.thoughtcrime.securesms.testing
+
+import okhttp3.mockwebserver.MockResponse
+import okhttp3.mockwebserver.RecordedRequest
+import okhttp3.mockwebserver.SocketPolicy
+import org.thoughtcrime.securesms.util.JsonUtils
+import java.util.concurrent.TimeUnit
+
+typealias ResponseFactory = (request: RecordedRequest) -> MockResponse
+
+/**
+ * Represent an HTTP verb for mocking web requests.
+ */
+sealed class Verb(val verb: String, val path: String, val responseFactory: ResponseFactory)
+
+class Get(path: String, responseFactory: ResponseFactory) : Verb("GET", path, responseFactory)
+
+class Put(path: String, responseFactory: ResponseFactory) : Verb("PUT", path, responseFactory)
+
+fun MockResponse.success(response: Any? = null): MockResponse {
+  return setResponseCode(200).apply {
+    if (response != null) {
+      setBody(JsonUtils.toJson(response))
+    }
+  }
+}
+
+fun MockResponse.failure(code: Int, response: Any? = null): MockResponse {
+  return setResponseCode(code).apply {
+    if (response != null) {
+      setBody(JsonUtils.toJson(response))
+    }
+  }
+}
+
+fun MockResponse.connectionFailure(): MockResponse {
+  return setSocketPolicy(SocketPolicy.DISCONNECT_AT_START)
+}
+
+fun MockResponse.timeout(): MockResponse {
+  return setHeadersDelay(1, TimeUnit.DAYS)
+    .setBodyDelay(1, TimeUnit.DAYS)
+}
+
+inline fun <reified T> RecordedRequest.parsedRequestBody(): T {
+  val bodyString = String(body.readByteArray())
+  return JsonUtils.fromJson(bodyString, T::class.java)
+}

app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalActivityRule.kt

@@ -8,6 +8,7 @@ import android.content.SharedPreferences
 import android.preference.PreferenceManager
 import androidx.test.core.app.ActivityScenario
 import androidx.test.platform.app.InstrumentationRegistry
+import okhttp3.mockwebserver.MockResponse
 import org.junit.rules.ExternalResource
 import org.signal.libsignal.protocol.IdentityKey
 import org.signal.libsignal.protocol.SignalProtocolAddress
@@ -17,6 +18,7 @@ import org.thoughtcrime.securesms.crypto.ProfileKeyUtil
 import org.thoughtcrime.securesms.database.IdentityDatabase
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.dependencies.ApplicationDependencies
+import org.thoughtcrime.securesms.dependencies.InstrumentationApplicationDependencyProvider
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.net.DeviceTransferBlockingInterceptor
 import org.thoughtcrime.securesms.profiles.ProfileName
@@ -29,6 +31,8 @@ import org.thoughtcrime.securesms.util.Util
 import org.whispersystems.signalservice.api.profiles.SignalServiceProfile
 import org.whispersystems.signalservice.api.push.ACI
 import org.whispersystems.signalservice.api.push.SignalServiceAddress
+import org.whispersystems.signalservice.internal.ServiceResponse
+import org.whispersystems.signalservice.internal.ServiceResponseProcessor
 import org.whispersystems.signalservice.internal.push.VerifyAccountResponse
 import java.lang.IllegalArgumentException
 import java.util.UUID
@@ -54,6 +58,8 @@ class SignalActivityRule(private val othersCount: Int = 4) : ExternalResource()
     context = InstrumentationRegistry.getInstrumentation().targetContext
     self = setupSelf()
     others = setupOthers()
+
+    InstrumentationApplicationDependencyProvider.clearHandlers()
   }
 
   private fun setupSelf(): Recipient {
@@ -67,18 +73,22 @@ class SignalActivityRule(private val othersCount: Int = 4) : ExternalResource()
 
     val registrationRepository = RegistrationRepository(application)
 
-    registrationRepository.registerAccountWithoutRegistrationLock(
+    InstrumentationApplicationDependencyProvider.addMockWebRequestHandlers(Put("/v2/keys") { MockResponse().success() })
+    val response: ServiceResponse<VerifyAccountResponse> = registrationRepository.registerAccountWithoutRegistrationLock(
       RegistrationData(
         code = "123123",
-        e164 = "+15554045550101",
+        e164 = "+15555550101",
         password = Util.getSecret(18),
         registrationId = registrationRepository.registrationId,
-        profileKey = registrationRepository.getProfileKey("+15554045550101"),
-        fcmToken = null
+        profileKey = registrationRepository.getProfileKey("+15555550101"),
+        fcmToken = null,
+        pniRegistrationId = registrationRepository.pniRegistrationId
       ),
       VerifyAccountResponse(UUID.randomUUID().toString(), UUID.randomUUID().toString(), false)
     ).blockingGet()
 
+    ServiceResponseProcessor.DefaultProcessor(response).resultOrThrow
+
     SignalStore.kbsValues().optOut()
     RegistrationUtil.maybeMarkRegistrationComplete(application)
     SignalDatabase.recipients.setProfileName(Recipient.self().id, ProfileName.fromParts("Tester", "McTesterson"))
@@ -107,7 +117,7 @@ class SignalActivityRule(private val othersCount: Int = 4) : ExternalResource()
     return others
   }
 
-  inline fun <reified T : Activity> launchActivity(initIntent: Intent.() -> Unit): ActivityScenario<T> {
+  inline fun <reified T : Activity> launchActivity(initIntent: Intent.() -> Unit = {}): ActivityScenario<T> {
     return androidx.test.core.app.launchActivity(Intent(context, T::class.java).apply(initIntent))
   }
 

app/src/androidTest/java/org/thoughtcrime/securesms/testing/SignalTestRunner.kt

@@ -0,0 +1,16 @@
+package org.thoughtcrime.securesms.testing
+
+import android.app.Application
+import android.content.Context
+import androidx.test.runner.AndroidJUnitRunner
+import org.thoughtcrime.securesms.SignalInstrumentationApplicationContext
+
+/**
+ * Custom runner that replaces application with [SignalInstrumentationApplicationContext].
+ */
+@Suppress("unused")
+class SignalTestRunner : AndroidJUnitRunner() {
+  override fun newApplication(cl: ClassLoader?, className: String?, context: Context?): Application {
+    return super.newApplication(cl, SignalInstrumentationApplicationContext::class.java.name, context)
+  }
+}

app/src/androidTest/java/org/thoughtcrime/securesms/testing/TestUtils.kt

@@ -0,0 +1,46 @@
+package org.thoughtcrime.securesms.testing
+
+import org.hamcrest.MatcherAssert.assertThat
+import org.hamcrest.Matchers.hasSize
+import org.hamcrest.Matchers.`is`
+import org.hamcrest.Matchers.not
+import org.hamcrest.Matchers.notNullValue
+import org.hamcrest.Matchers.nullValue
+import java.util.concurrent.CountDownLatch
+
+/**
+ * Run the given [runnable] on a new thread and wait for it to finish.
+ */
+fun runSync(runnable: () -> Unit) {
+  val lock = CountDownLatch(1)
+  Thread {
+    try {
+      runnable.invoke()
+    } finally {
+      lock.countDown()
+    }
+  }.start()
+  lock.await()
+}
+
+/* Various kotlin-ifications of hamcrest matchers */
+
+fun <T : Any?> T.assertIsNull() {
+  assertThat(this, nullValue())
+}
+
+fun <T : Any?> T.assertIsNotNull() {
+  assertThat(this, notNullValue())
+}
+
+infix fun <T : Any> T.assertIs(expected: T) {
+  assertThat(this, `is`(expected))
+}
+
+infix fun <T : Any> T.assertIsNot(expected: T) {
+  assertThat(this, not(`is`(expected)))
+}
+
+infix fun <E, T : Collection<E>> T.assertIsSize(expected: Int) {
+  assertThat(this, hasSize(expected))
+}