From 905b0681f559d22ded062a33cba02715606c957f Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Wed, 22 May 2024 16:00:48 -0400
Subject: [PATCH] Update otpk/kpk tests.

---
 .../securesms/database/KyberPreKeyTableTest.kt        | 11 +++++++++--
 .../securesms/database/OneTimePreKeyTableTest.kt      | 11 +++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt
index 4736aac117..e52deb27ec 100644
--- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/database/KyberPreKeyTableTest.kt
@@ -159,7 +159,7 @@ class KyberPreKeyTableTest {
     val count = SignalDatabase.rawDatabase
       .update(KyberPreKeyTable.TABLE_NAME)
       .values(KyberPreKeyTable.STALE_TIMESTAMP to staleTime)
-      .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account)
+      .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId())
       .run()
 
     assertEquals(1, count)
@@ -169,8 +169,15 @@ class KyberPreKeyTableTest {
     return SignalDatabase.rawDatabase
       .select(KyberPreKeyTable.STALE_TIMESTAMP)
       .from(KyberPreKeyTable.TABLE_NAME)
-      .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account)
+      .where("${KyberPreKeyTable.ACCOUNT_ID} = ? AND ${KyberPreKeyTable.KEY_ID} = $id", account.toAccountId())
       .run()
       .readToSingleObject { it.requireLongOrNull(KyberPreKeyTable.STALE_TIMESTAMP) }
   }
+
+  private fun ServiceId.toAccountId(): String {
+    return when (this) {
+      is ACI -> this.toString()
+      is PNI -> KyberPreKeyTable.PNI_ACCOUNT_ID
+    }
+  }
 }
diff --git a/app/src/androidTest/java/org/thoughtcrime/securesms/database/OneTimePreKeyTableTest.kt b/app/src/androidTest/java/org/thoughtcrime/securesms/database/OneTimePreKeyTableTest.kt
index beaacf0d34..e28b54b891 100644
--- a/app/src/androidTest/java/org/thoughtcrime/securesms/database/OneTimePreKeyTableTest.kt
+++ b/app/src/androidTest/java/org/thoughtcrime/securesms/database/OneTimePreKeyTableTest.kt
@@ -120,7 +120,7 @@ class OneTimePreKeyTableTest {
     val count = SignalDatabase.rawDatabase
       .update(OneTimePreKeyTable.TABLE_NAME)
       .values(OneTimePreKeyTable.STALE_TIMESTAMP to staleTime)
-      .where("${OneTimePreKeyTable.ACCOUNT_ID} = ? AND ${OneTimePreKeyTable.KEY_ID} = $id", account)
+      .where("${OneTimePreKeyTable.ACCOUNT_ID} = ? AND ${OneTimePreKeyTable.KEY_ID} = $id", account.toAccountId())
       .run()
 
     assertEquals(1, count)
@@ -130,8 +130,15 @@ class OneTimePreKeyTableTest {
     return SignalDatabase.rawDatabase
       .select(OneTimePreKeyTable.STALE_TIMESTAMP)
       .from(OneTimePreKeyTable.TABLE_NAME)
-      .where("${OneTimePreKeyTable.ACCOUNT_ID} = ? AND ${OneTimePreKeyTable.KEY_ID} = $id", account)
+      .where("${OneTimePreKeyTable.ACCOUNT_ID} = ? AND ${OneTimePreKeyTable.KEY_ID} = $id", account.toAccountId())
       .run()
       .readToSingleObject { it.requireLongOrNull(OneTimePreKeyTable.STALE_TIMESTAMP) }
   }
+
+  private fun ServiceId.toAccountId(): String {
+    return when (this) {
+      is ACI -> this.toString()
+      is PNI -> OneTimePreKeyTable.PNI_ACCOUNT_ID
+    }
+  }
 }