From 5e9824a1804f88b4c72d8ca0f1eea4abf1bed760 Mon Sep 17 00:00:00 2001
From: Alex Bakon <akonradi@signal.org>
Date: Thu, 15 May 2025 10:26:00 -0400
Subject: [PATCH] Update test client prekey usage to include kyber keys.

---
 .../thoughtcrime/securesms/testing/BobClient.kt   | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/testing/BobClient.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/BobClient.kt
index a648647344..237f176d9c 100644
--- a/app/src/androidTest/java/org/thoughtcrime/securesms/testing/BobClient.kt
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/testing/BobClient.kt
@@ -19,6 +19,7 @@ import org.signal.libsignal.protocol.util.KeyHelper
 import org.signal.libsignal.zkgroup.profiles.ProfileKey
 import org.thoughtcrime.securesms.crypto.ProfileKeyUtil
 import org.thoughtcrime.securesms.crypto.SealedSenderAccessUtil
+import org.thoughtcrime.securesms.database.KyberPreKeyTable
 import org.thoughtcrime.securesms.database.OneTimePreKeyTable
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.database.SignedPreKeyTable
@@ -102,6 +103,15 @@ class BobClient(val serviceId: ServiceId, val e164: String, val identityKeyPair:
 
     val selfSignedPreKeyRecord = SignalDatabase.signedPreKeys.get(getAliceServiceId(), selfSignedPreKeyId)!!
 
+    val selfSignedKyberPreKeyId = SignalDatabase.rawDatabase
+      .select(KyberPreKeyTable.KEY_ID)
+      .from(KyberPreKeyTable.TABLE_NAME)
+      .where("${KyberPreKeyTable.ACCOUNT_ID} = ?", getAliceServiceId().toString())
+      .run()
+      .readToSingleInt(-1)
+
+    val selfSignedKyberPreKeyRecord = SignalDatabase.kyberPreKeys.get(getAliceServiceId(), selfSignedKyberPreKeyId)!!.record
+
     return PreKeyBundle(
       SignalStore.account.registrationId,
       1,
@@ -110,7 +120,10 @@ class BobClient(val serviceId: ServiceId, val e164: String, val identityKeyPair:
       selfSignedPreKeyId,
       selfSignedPreKeyRecord.keyPair.publicKey,
       selfSignedPreKeyRecord.signature,
-      getAlicePublicKey()
+      getAlicePublicKey(),
+      selfSignedKyberPreKeyId,
+      selfSignedKyberPreKeyRecord.keyPair.publicKey,
+      selfSignedKyberPreKeyRecord.signature
     )
   }