From 315e0736f5970b0177087183496319cfa411140c Mon Sep 17 00:00:00 2001
From: Steve Soltys <steve@stevesoltys.com>
Date: Wed, 25 Sep 2024 02:38:12 +0000
Subject: [PATCH] Fix KVRestore spies to use new functions

---
 .../seedvault/e2e/LargeRestoreTestBase.kt     | 42 +++++++++++++++----
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/app/src/androidTest/java/com/stevesoltys/seedvault/e2e/LargeRestoreTestBase.kt b/app/src/androidTest/java/com/stevesoltys/seedvault/e2e/LargeRestoreTestBase.kt
index a2384191..d104d3ed 100644
--- a/app/src/androidTest/java/com/stevesoltys/seedvault/e2e/LargeRestoreTestBase.kt
+++ b/app/src/androidTest/java/com/stevesoltys/seedvault/e2e/LargeRestoreTestBase.kt
@@ -14,6 +14,8 @@ import com.stevesoltys.seedvault.e2e.screen.impl.RestoreScreen
 import com.stevesoltys.seedvault.transport.restore.FullRestore
 import com.stevesoltys.seedvault.transport.restore.KVRestore
 import com.stevesoltys.seedvault.transport.restore.OutputFactory
+import io.mockk.Call
+import io.mockk.MockKAnswerScope
 import io.mockk.clearMocks
 import io.mockk.coEvery
 import io.mockk.every
@@ -161,14 +163,26 @@ internal interface LargeRestoreTestBase : LargeTestBase {
 
         clearMocks(spyKVRestore)
 
-        coEvery {
-            spyKVRestore.initializeState(any(), any(), any(), any())
-        } answers {
-            packageName = arg<PackageInfo>(3).packageName
+        fun initializeStateBlock(
+            packageInfoIndex: Int
+        ): MockKAnswerScope<Unit, Unit>.(Call) -> Unit = {
+            packageName = arg<PackageInfo>(packageInfoIndex).packageName
             restoreResult.kv[packageName!!] = mutableMapOf()
             callOriginal()
         }
 
+        coEvery {
+            spyKVRestore.initializeState(any(), any(), any(), any())
+        } answers initializeStateBlock(1)
+
+        coEvery {
+            spyKVRestore.initializeStateV1(any(), any(), any(), any())
+        } answers initializeStateBlock(2)
+
+        coEvery {
+            spyKVRestore.initializeStateV0(any(), any())
+        } answers initializeStateBlock(1)
+
         every {
             spyOutputFactory.getBackupDataOutput(any())
         } answers {
@@ -186,19 +200,31 @@ internal interface LargeRestoreTestBase : LargeTestBase {
 
         clearMocks(spyFullRestore)
 
-        coEvery {
-            spyFullRestore.initializeState(any(), any(), any())
-        } answers {
+        fun initializeStateBlock(
+            packageInfoIndex: Int
+        ): MockKAnswerScope<Unit, Unit>.(Call) -> Unit = {
             packageName?.let {
                 restoreResult.full[it] = dataIntercept.toByteArray().sha256()
             }
 
-            packageName = arg<PackageInfo>(3).packageName
+            packageName = arg<PackageInfo>(packageInfoIndex).packageName
             dataIntercept = ByteArrayOutputStream()
 
             callOriginal()
         }
 
+        coEvery {
+            spyFullRestore.initializeState(any(), any(), any())
+        } answers initializeStateBlock(1)
+
+        coEvery {
+            spyFullRestore.initializeStateV1(any(), any(), any())
+        } answers initializeStateBlock(2)
+
+        coEvery {
+            spyFullRestore.initializeStateV0(any(), any())
+        } answers initializeStateBlock(1)
+
         every {
             spyOutputFactory.getOutputStream(any())
         } answers {