From 69e7dfd90a81fbcd7667d266bdffda80cd6dcd1a Mon Sep 17 00:00:00 2001
From: Adam Brown <adampsbrown@gmail.com>
Date: Tue, 27 Sep 2022 21:13:21 +0100
Subject: [PATCH] using proxy to allow the encryption tests to wait for keys to
 be uploaded, re-enabling the 2nd device user tests

---
 test-harness/src/test/kotlin/SmokeTest.kt     | 26 +++++++++---------
 .../src/test/kotlin/test/TestMatrix.kt        | 27 +++++++++++++++++--
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/test-harness/src/test/kotlin/SmokeTest.kt b/test-harness/src/test/kotlin/SmokeTest.kt
index cccb5c8..488680c 100644
--- a/test-harness/src/test/kotlin/SmokeTest.kt
+++ b/test-harness/src/test/kotlin/SmokeTest.kt
@@ -123,18 +123,20 @@ class SmokeTest {
         alice.expectTextMessage(SharedState.sharedRoom, message2)
 
         // Needs investigation
-//        val aliceSecondDevice = testMatrix(SharedState.alice, isTemp = true, withLogging = true).also { it.newlogin() }
-//        aliceSecondDevice.client.syncService().startSyncing().collectAsync {
-//            val message3 = "from alice to bob and alice's second device".from(SharedState.alice.roomMember)
-//            alice.sendTextMessage(SharedState.sharedRoom, message3.content, isEncrypted)
-//            aliceSecondDevice.expectTextMessage(SharedState.sharedRoom, message3)
-//            bob.expectTextMessage(SharedState.sharedRoom, message3)
-//
-//            val message4 = "from alice's second device to bob and alice's first device".from(SharedState.alice.roomMember)
-//            aliceSecondDevice.sendTextMessage(SharedState.sharedRoom, message4.content, isEncrypted)
-//            alice.expectTextMessage(SharedState.sharedRoom, message4)
-//            bob.expectTextMessage(SharedState.sharedRoom, message4)
-//        }
+        val aliceSecondDevice = testMatrix(SharedState.alice, isTemp = true, withLogging = true).also { it.newlogin() }
+        aliceSecondDevice.client.syncService().startSyncing().collectAsync {
+            aliceSecondDevice.client.proxyDeviceService().waitForOneTimeKeysToBeUploaded()
+
+            val message3 = "from alice to bob and alice's second device".from(SharedState.alice.roomMember)
+            alice.sendTextMessage(SharedState.sharedRoom, message3.content, isEncrypted)
+            aliceSecondDevice.expectTextMessage(SharedState.sharedRoom, message3)
+            bob.expectTextMessage(SharedState.sharedRoom, message3)
+
+            val message4 = "from alice's second device to bob and alice's first device".from(SharedState.alice.roomMember)
+            aliceSecondDevice.sendTextMessage(SharedState.sharedRoom, message4.content, isEncrypted)
+            alice.expectTextMessage(SharedState.sharedRoom, message4)
+            bob.expectTextMessage(SharedState.sharedRoom, message4)
+        }
     }
 }
 
diff --git a/test-harness/src/test/kotlin/test/TestMatrix.kt b/test-harness/src/test/kotlin/test/TestMatrix.kt
index 023a3b1..fce9c4f 100644
--- a/test-harness/src/test/kotlin/test/TestMatrix.kt
+++ b/test-harness/src/test/kotlin/test/TestMatrix.kt
@@ -14,6 +14,7 @@ import app.dapk.st.matrix.crypto.RoomMembersProvider
 import app.dapk.st.matrix.crypto.Verification
 import app.dapk.st.matrix.crypto.cryptoService
 import app.dapk.st.matrix.crypto.installCryptoService
+import app.dapk.st.matrix.device.DeviceService
 import app.dapk.st.matrix.device.deviceService
 import app.dapk.st.matrix.device.installEncryptionService
 import app.dapk.st.matrix.http.ktor.KtorMatrixHttpClientFactory
@@ -39,6 +40,7 @@ import test.impl.PrintingErrorTracking
 import java.io.File
 import java.time.Clock
 import javax.imageio.ImageIO
+import kotlin.coroutines.resume
 
 object TestUsers {
 
@@ -93,7 +95,9 @@ class TestMatrix(
     ).also {
         it.install {
             installAuthService(storeModule.credentialsStore())
-            installEncryptionService(storeModule.knownDevicesStore())
+            installEncryptionService(storeModule.knownDevicesStore()).proxy {
+                ProxyDeviceService(it)
+            }
 
             val olmAccountStore = OlmPersistenceWrapper(storeModule.olmStore(), base64)
             val olm = OlmWrapper(
@@ -355,4 +359,23 @@ class JavaImageContentReader : ImageContentReader {
 
     override fun inputStream(uri: String) = File(uri).inputStream()
 
-}
\ No newline at end of file
+}
+
+class ProxyDeviceService(private val deviceService: DeviceService) : DeviceService by deviceService {
+
+    private var oneTimeKeysContinuation: (() -> Unit)? = null
+
+    override suspend fun uploadOneTimeKeys(oneTimeKeys: DeviceService.OneTimeKeys) {
+        deviceService.uploadOneTimeKeys(oneTimeKeys)
+        oneTimeKeysContinuation?.invoke()?.also { oneTimeKeysContinuation = null }
+    }
+
+    suspend fun waitForOneTimeKeysToBeUploaded() {
+        suspendCancellableCoroutine { continuation ->
+            oneTimeKeysContinuation = { continuation.resume(Unit) }
+        }
+    }
+
+}
+
+fun MatrixClient.proxyDeviceService() = this.deviceService() as ProxyDeviceService
\ No newline at end of file