1 files changed, 43 insertions, 35 deletions

diff --git a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
index 2bd9b49e..4b022e5e 100644
--- a/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
+++ b/crypto-pgpainless/src/test/kotlin/dev/msfjarvis/aps/crypto/PGPKeyManagerTest.kt
@@ -3,32 +3,48 @@ package dev.msfjarvis.aps.crypto
 import com.github.michaelbull.result.unwrap
 import com.github.michaelbull.result.unwrapError
 import java.io.File
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestCoroutineDispatcher
-import kotlinx.coroutines.test.runBlockingTest
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
 import org.junit.Rule
-import org.junit.Test
 import org.junit.rules.TemporaryFolder
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class PGPKeyManagerTest {
 
   @get:Rule val temporaryFolder: TemporaryFolder = TemporaryFolder()
-  private val filesDir by lazy(LazyThreadSafetyMode.NONE) { temporaryFolder.root }
-  private val keysDir by
-    lazy(LazyThreadSafetyMode.NONE) { File(filesDir, PGPKeyManager.KEY_DIR_NAME) }
-  private val testCoroutineDispatcher = TestCoroutineDispatcher()
-  private val keyManager by
-    lazy(LazyThreadSafetyMode.NONE) {
-      PGPKeyManager(filesDir.absolutePath, testCoroutineDispatcher)
-    }
+  private val filesDir by unsafeLazy { temporaryFolder.root }
+  private val keysDir by unsafeLazy { File(filesDir, PGPKeyManager.KEY_DIR_NAME) }
+  private val dispatcher = StandardTestDispatcher()
+  private val scope = TestScope(dispatcher)
+  private val keyManager by unsafeLazy { PGPKeyManager(filesDir.absolutePath, dispatcher) }
   private val key = PGPKeyManager.makeKey(TestUtils.getArmoredPrivateKey())
 
+  private fun <T> unsafeLazy(initializer: () -> T) =
+    lazy(LazyThreadSafetyMode.NONE) { initializer.invoke() }
+
+  @BeforeTest
+  fun setUp() {
+    Dispatchers.setMain(dispatcher)
+  }
+
+  @AfterTest
+  fun tearDown() {
+    Dispatchers.resetMain()
+  }
+
   @Test
-  fun testAddingKey() {
-    runBlockingTest {
+  fun testAddingKey() =
+    scope.runTest {
       // Check if the key id returned is correct
       val keyId = keyManager.addKey(key).unwrap().getKeyId()
       assertEquals(CryptoConstants.KEY_ID, keyId)
@@ -40,33 +56,30 @@ class PGPKeyManagerTest {
       val keyFile = keysDir.listFiles()?.first()
       assertEquals(keyFile?.name, "$keyId.${PGPKeyManager.KEY_EXTENSION}")
     }
-  }
 
   @Test
-  fun testAddingKeyWithoutReplaceFlag() {
-    runBlockingTest {
+  fun testAddingKeyWithoutReplaceFlag() =
+    scope.runTest {
       // Check adding the keys twice
       keyManager.addKey(key, false).unwrap()
       val error = keyManager.addKey(key, false).unwrapError()
 
       assertIs<KeyManagerException.KeyAlreadyExistsException>(error)
     }
-  }
 
   @Test
-  fun testAddingKeyWithReplaceFlag() {
-    runBlockingTest {
+  fun testAddingKeyWithReplaceFlag() =
+    scope.runTest {
       // Check adding the keys twice
       keyManager.addKey(key, true).unwrap()
       val keyId = keyManager.addKey(key, true).unwrap().getKeyId()
 
       assertEquals(CryptoConstants.KEY_ID, keyId)
     }
-  }
 
   @Test
-  fun testRemovingKey() {
-    runBlockingTest {
+  fun testRemovingKey() =
+    scope.runTest {
       // Add key using KeyManager
       keyManager.addKey(key).unwrap()
 
@@ -78,11 +91,10 @@ class PGPKeyManagerTest {
       val keysDir = File(filesDir, PGPKeyManager.KEY_DIR_NAME)
       assertEquals(0, keysDir.list()?.size)
     }
-  }
 
   @Test
-  fun testGetExistingKey() {
-    runBlockingTest {
+  fun testGetExistingKey() =
+    scope.runTest {
       // Add key using KeyManager
       keyManager.addKey(key).unwrap()
 
@@ -91,11 +103,10 @@ class PGPKeyManagerTest {
       assertEquals(CryptoConstants.KEY_ID, key.getKeyId())
       assertEquals(key.getKeyId(), returnedKeyPair.getKeyId())
     }
-  }
 
   @Test
-  fun testGetNonExistentKey() {
-    runBlockingTest {
+  fun testGetNonExistentKey() =
+    scope.runTest {
       // Add key using KeyManager
       keyManager.addKey(key).unwrap()
 
@@ -106,21 +117,19 @@ class PGPKeyManagerTest {
       assertIs<KeyManagerException.KeyNotFoundException>(error)
       assertEquals("No key found with id: $randomKeyId", error.message)
     }
-  }
 
   @Test
-  fun testFindKeysWithoutAdding() {
-    runBlockingTest {
+  fun testFindKeysWithoutAdding() =
+    scope.runTest {
       // Check returned key
       val error = keyManager.getKeyById("0x123456789").unwrapError()
       assertIs<KeyManagerException.NoKeysAvailableException>(error)
       assertEquals("No keys were found", error.message)
     }
-  }
 
   @Test
-  fun testGettingAllKeys() {
-    runBlockingTest {
+  fun testGettingAllKeys() =
+    scope.runTest {
       // TODO: Should we check for more than 1 keys?
       // Check if KeyManager returns no key
       val noKeyList = keyManager.getAllKeys().unwrap()
@@ -133,5 +142,4 @@ class PGPKeyManagerTest {
       val singleKeyList = keyManager.getAllKeys().unwrap()
       assertEquals(1, singleKeyList.size)
     }
-  }
 }