public Account updatePniKeys(final Account account,
      final IdentityKey pniIdentityKey,
      final Map<Byte, ECSignedPreKey> deviceSignedPreKeys,
      @Nullable final Map<Byte, KEMSignedPreKey> devicePqLastResortPreKeys,
      final List<IncomingMessage> deviceMessages,
      final Map<Byte, Integer> pniRegistrationIds) throws MismatchedDevicesException, StaleDevicesException {
    validateDeviceMessages(account, deviceMessages);

    // Don't try to be smart about ignoring unnecessary retries. If we make literally no change we will skip the ddb
    // write anyway. Linked devices can handle some wasted extra key rotations.
    final Account updatedAccount = accountsManager.updatePniKeys(
        account, pniIdentityKey, deviceSignedPreKeys, devicePqLastResortPreKeys, pniRegistrationIds);

    sendDeviceMessages(updatedAccount, deviceMessages);
    return updatedAccount;
  }
@Test
  void updatePniKeysSetPrimaryDevicePrekeyAndSendMessages() throws Exception {
    final UUID aci = UUID.randomUUID();
    final UUID pni = UUID.randomUUID();

    final Account account = mock(Account.class);
    when(account.getUuid()).thenReturn(aci);
    when(account.getPhoneNumberIdentifier()).thenReturn(pni);

    final Device d2 = mock(Device.class);
    final byte deviceId2 = 2;
    when(d2.getId()).thenReturn(deviceId2);

    when(account.getDevice(deviceId2)).thenReturn(Optional.of(d2));
    when(account.getDevices()).thenReturn(List.of(d2));

    final ECKeyPair pniIdentityKeyPair = Curve.generateKeyPair();
    final IdentityKey pniIdentityKey = new IdentityKey(pniIdentityKeyPair.getPublicKey());
    final Map<Byte, ECSignedPreKey> prekeys = Map.of(Device.PRIMARY_ID,
        KeysHelper.signedECPreKey(1, pniIdentityKeyPair),
        deviceId2, KeysHelper.signedECPreKey(2, pniIdentityKeyPair));
    final Map<Byte, Integer> registrationIds = Map.of(Device.PRIMARY_ID, 17, deviceId2, 19);

    final IncomingMessage msg = mock(IncomingMessage.class);
    when(msg.destinationDeviceId()).thenReturn(deviceId2);
    when(msg.content()).thenReturn(Base64.getEncoder().encodeToString(new byte[]{1}));

    changeNumberManager.updatePniKeys(account, pniIdentityKey, prekeys, null, List.of(msg), registrationIds);

    verify(accountsManager).updatePniKeys(account, pniIdentityKey, prekeys, null, registrationIds);

    final ArgumentCaptor<MessageProtos.Envelope> envelopeCaptor = ArgumentCaptor.forClass(MessageProtos.Envelope.class);
    verify(messageSender).sendMessage(any(), eq(d2), envelopeCaptor.capture(), eq(false));

    final MessageProtos.Envelope envelope = envelopeCaptor.getValue();

    assertEquals(aci, UUID.fromString(envelope.getDestinationUuid()));
    assertEquals(aci, UUID.fromString(envelope.getSourceUuid()));
    assertEquals(Device.PRIMARY_ID, envelope.getSourceDevice());
    assertFalse(updatedPhoneNumberIdentifiersByAccount.containsKey(account));
  }