PIR: Update completeRequestData to support scans and email extracted data

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/actions/EmailDataReceivedEventHandler.kt

@@ -42,6 +42,7 @@ class EmailDataReceivedEventHandler @Inject constructor() : EventHandler {
         return Next(
             nextState = state.copy(
                 currentActionIndex = state.currentActionIndex + 1,
+                actionRetryCount = 0,
                 emailExtractedData = actualEvent.emailExtractedData,
             ),
             nextEvent = ExecuteBrokerStepAction(

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/common/actions/ExecuteBrokerStepActionEventHandler.kt

@@ -40,7 +40,6 @@ import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.SideEf
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.SideEffect.PushJsAction
 import com.duckduckgo.pir.impl.common.actions.PirActionsRunnerStateEngine.State
 import com.duckduckgo.pir.impl.common.toParams
-import com.duckduckgo.pir.impl.models.ProfileQuery
 import com.duckduckgo.pir.impl.pixels.PirStage
 import com.duckduckgo.pir.impl.scripts.models.BrokerAction
 import com.duckduckgo.pir.impl.scripts.models.BrokerAction.Click
@@ -52,10 +51,10 @@ import com.duckduckgo.pir.impl.scripts.models.BrokerAction.GetCaptchaInfo
 import com.duckduckgo.pir.impl.scripts.models.BrokerAction.GetEmailData
 import com.duckduckgo.pir.impl.scripts.models.BrokerAction.SolveCaptcha
 import com.duckduckgo.pir.impl.scripts.models.DataSource.EXTRACTED_PROFILE
+import com.duckduckgo.pir.impl.scripts.models.ExtractedProfileParams
 import com.duckduckgo.pir.impl.scripts.models.PirError
 import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData
 import com.duckduckgo.pir.impl.scripts.models.PirScriptRequestData.UserProfile
-import com.duckduckgo.pir.impl.store.PirRepository.GeneratedEmailData
 import com.squareup.anvil.annotations.ContributesMultibinding
 import javax.inject.Inject
 import kotlin.reflect.KClass
@@ -283,7 +282,7 @@ class ExecuteBrokerStepActionEventHandler @Inject constructor(
                             actionToExecute.id,
                             actionToExecute,
                             pushDelay,
-                            completeRequestData(currentBrokerStep, actionToExecute, state.profileQuery, requestData, state.generatedEmailData),
+                            completeRequestData(currentBrokerStep, actionToExecute, state, requestData),
                         ),
                     )
                 }
@@ -294,35 +293,36 @@ class ExecuteBrokerStepActionEventHandler @Inject constructor(
     private fun completeRequestData(
         brokerStep: BrokerStep,
         actionToExecute: BrokerAction,
-        profileQuery: ProfileQuery,
+        state: State,
         requestData: PirScriptRequestData,
-        generatedEmailData: GeneratedEmailData?,
     ): PirScriptRequestData {
-        val extractedProfile = if (brokerStep is OptOutStep && actionToExecute.dataSource == EXTRACTED_PROFILE &&
-            (requestData as UserProfile).extractedProfile == null
-        ) {
-            brokerStep.profileToOptOut
-        } else if (brokerStep is EmailConfirmationStep && actionToExecute.dataSource == EXTRACTED_PROFILE &&
-            (requestData as UserProfile).extractedProfile == null
-        ) {
-            brokerStep.profileToOptOut
-        } else {
-            null
+        if (requestData !is UserProfile || requestData.extractedProfile != null) {
+            return requestData
+        }
+        if (actionToExecute.dataSource != EXTRACTED_PROFILE) {
+            return requestData
         }
 
-        return if (extractedProfile != null && requestData is UserProfile) {
-            val params = extractedProfile.toParams(profileQuery.fullName)
-            UserProfile(
-                userProfile = requestData.userProfile,
-                extractedProfile = if (generatedEmailData != null) {
-                    params.copy(email = generatedEmailData.emailAddress)
-                } else {
-                    params
-                },
-            )
+        val baseParams: ExtractedProfileParams = when (brokerStep) {
+            is OptOutStep -> brokerStep.profileToOptOut.toParams(state.profileQuery.fullName)
+            is EmailConfirmationStep -> brokerStep.profileToOptOut.toParams(state.profileQuery.fullName)
+            is ScanStep -> if (state.generatedEmailData != null) ExtractedProfileParams() else return requestData
+        }
+
+        val withEmail = state.generatedEmailData?.let {
+            baseParams.copy(email = it.emailAddress)
+        } ?: baseParams
+
+        val withEmailExtractedData = if (state.emailExtractedData.isNotEmpty()) {
+            withEmail.copy(emailExtractedData = state.emailExtractedData)
         } else {
-            requestData
+            withEmail
         }
+
+        return UserProfile(
+            userProfile = requestData.userProfile,
+            extractedProfile = withEmailExtractedData,
+        )
     }
 
     companion object {

pir/pir-impl/src/main/java/com/duckduckgo/pir/impl/scripts/models/PirScriptRequestParams.kt

@@ -44,4 +44,5 @@ data class ExtractedProfileParams(
     val profileUrl: String? = null,
     val email: String? = null,
     val fullName: String? = null,
+    val emailExtractedData: Map<String, String>? = null,
 )

pir/pir-impl/src/test/kotlin/com/duckduckgo/pir/impl/common/actions/EmailDataReceivedEventHandlerTest.kt

@@ -124,15 +124,23 @@ class EmailDataReceivedEventHandlerTest {
     fun whenEmailDataReceivedThenPreservesOtherStateFields() = runTest {
         val state = baseState.copy(
             currentBrokerStepIndex = 4,
-            actionRetryCount = 1,
             attemptId = "test-attempt",
         )
         val event = EmailDataReceived(emailExtractedData = mapOf("verificationCode" to "abc"))
 
         val result = testee.invoke(state, event)
 
         assertEquals(4, result.nextState.currentBrokerStepIndex)
-        assertEquals(1, result.nextState.actionRetryCount)
         assertEquals("test-attempt", result.nextState.attemptId)
     }
+
+    @Test
+    fun whenEmailDataReceivedThenResetsActionRetryCount() = runTest {
+        val state = baseState.copy(actionRetryCount = 3)
+        val event = EmailDataReceived(emailExtractedData = mapOf("verificationCode" to "abc"))
+
+        val result = testee.invoke(state, event)
+
+        assertEquals(0, result.nextState.actionRetryCount)
+    }
 }

pir/pir-impl/src/test/kotlin/com/duckduckgo/pir/impl/common/actions/ExecuteBrokerStepActionEventHandlerTest.kt

@@ -1061,4 +1061,203 @@ class ExecuteBrokerStepActionEventHandlerTest {
         val sideEffect = result.sideEffect as AwaitEmailData
         assertEquals("", sideEffect.emailAddress)
     }
+
+    @Test
+    fun whenScanStepFillFormWithGeneratedEmailDataThenIncludesEmailInRequestData() = runTest {
+        val action = BrokerAction.FillForm(
+            id = "action-fill",
+            elements = emptyList(),
+            selector = "form",
+            dataSource = DataSource.EXTRACTED_PROFILE,
+        )
+        val scanStep = ScanStep(
+            broker = testBroker,
+            step = ScanStepActions(
+                stepType = "scan",
+                actions = listOf(action),
+                scanType = "initial",
+            ),
+        )
+        val state = State(
+            runType = RunType.MANUAL,
+            brokerStepsToExecute = listOf(scanStep),
+            profileQuery = testProfileQuery,
+            currentBrokerStepIndex = 0,
+            currentActionIndex = 0,
+            generatedEmailData = GeneratedEmailData(
+                emailAddress = "scan-generated@example.com",
+                pattern = "pattern-123",
+            ),
+            stageStatus = PirStageStatus(
+                currentStage = PirStage.OTHER,
+                stageStartMs = 0,
+            ),
+        )
+        val event = ExecuteBrokerStepAction(UserProfile(userProfile = testProfileQuery))
+
+        val result = testee.invoke(state, event)
+
+        val sideEffect = result.sideEffect as PushJsAction
+        val userData = sideEffect.requestParamsData as UserProfile
+        assertEquals(testProfileQuery, userData.userProfile)
+        assertEquals("scan-generated@example.com", userData.extractedProfile?.email)
+        assertNull(userData.extractedProfile?.name)
+        assertNull(userData.extractedProfile?.emailExtractedData)
+    }
+
+    @Test
+    fun whenScanStepFillFormWithoutGeneratedEmailDataThenExtractedProfileIsNull() = runTest {
+        val action = BrokerAction.FillForm(
+            id = "action-fill",
+            elements = emptyList(),
+            selector = "form",
+            dataSource = DataSource.EXTRACTED_PROFILE,
+        )
+        val scanStep = ScanStep(
+            broker = testBroker,
+            step = ScanStepActions(
+                stepType = "scan",
+                actions = listOf(action),
+                scanType = "initial",
+            ),
+        )
+        val state = State(
+            runType = RunType.MANUAL,
+            brokerStepsToExecute = listOf(scanStep),
+            profileQuery = testProfileQuery,
+            currentBrokerStepIndex = 0,
+            currentActionIndex = 0,
+            generatedEmailData = null,
+            stageStatus = PirStageStatus(
+                currentStage = PirStage.OTHER,
+                stageStartMs = 0,
+            ),
+        )
+        val event = ExecuteBrokerStepAction(UserProfile(userProfile = testProfileQuery))
+
+        val result = testee.invoke(state, event)
+
+        val sideEffect = result.sideEffect as PushJsAction
+        val userData = sideEffect.requestParamsData as UserProfile
+        assertNull(userData.extractedProfile)
+    }
+
+    @Test
+    fun whenOptOutFillFormWithEmailExtractedDataThenIncludesItInRequestData() = runTest {
+        val action = BrokerAction.FillForm(
+            id = "action-fill",
+            elements = emptyList(),
+            selector = "form",
+            dataSource = DataSource.EXTRACTED_PROFILE,
+        )
+        val optOutStep = OptOutStep(
+            broker = testBroker,
+            step = OptOutStepActions(
+                stepType = "optout",
+                actions = listOf(action),
+                optOutType = "form",
+            ),
+            profileToOptOut = testExtractedProfile,
+        )
+        val state = State(
+            runType = RunType.OPTOUT,
+            brokerStepsToExecute = listOf(optOutStep),
+            profileQuery = testProfileQuery,
+            currentBrokerStepIndex = 0,
+            currentActionIndex = 0,
+            emailExtractedData = mapOf("verificationCode" to "123456"),
+            stageStatus = PirStageStatus(
+                currentStage = PirStage.OTHER,
+                stageStartMs = 0,
+            ),
+        )
+        val event = ExecuteBrokerStepAction(UserProfile(userProfile = testProfileQuery))
+
+        val result = testee.invoke(state, event)
+
+        val sideEffect = result.sideEffect as PushJsAction
+        val userData = sideEffect.requestParamsData as UserProfile
+        assertEquals("John Doe", userData.extractedProfile?.name)
+        assertEquals(mapOf("verificationCode" to "123456"), userData.extractedProfile?.emailExtractedData)
+    }
+
+    @Test
+    fun whenScanStepFillFormWithGeneratedEmailAndEmailExtractedDataThenIncludesBoth() = runTest {
+        val action = BrokerAction.FillForm(
+            id = "action-fill",
+            elements = emptyList(),
+            selector = "form",
+            dataSource = DataSource.EXTRACTED_PROFILE,
+        )
+        val scanStep = ScanStep(
+            broker = testBroker,
+            step = ScanStepActions(
+                stepType = "scan",
+                actions = listOf(action),
+                scanType = "initial",
+            ),
+        )
+        val state = State(
+            runType = RunType.MANUAL,
+            brokerStepsToExecute = listOf(scanStep),
+            profileQuery = testProfileQuery,
+            currentBrokerStepIndex = 0,
+            currentActionIndex = 0,
+            generatedEmailData = GeneratedEmailData(
+                emailAddress = "scan-generated@example.com",
+                pattern = "pattern-123",
+            ),
+            emailExtractedData = mapOf("verificationCode" to "654321"),
+            stageStatus = PirStageStatus(
+                currentStage = PirStage.OTHER,
+                stageStartMs = 0,
+            ),
+        )
+        val event = ExecuteBrokerStepAction(UserProfile(userProfile = testProfileQuery))
+
+        val result = testee.invoke(state, event)
+
+        val sideEffect = result.sideEffect as PushJsAction
+        val userData = sideEffect.requestParamsData as UserProfile
+        assertEquals("scan-generated@example.com", userData.extractedProfile?.email)
+        assertEquals(mapOf("verificationCode" to "654321"), userData.extractedProfile?.emailExtractedData)
+    }
+
+    @Test
+    fun whenEmailExtractedDataIsEmptyThenFieldIsNull() = runTest {
+        val action = BrokerAction.FillForm(
+            id = "action-fill",
+            elements = emptyList(),
+            selector = "form",
+            dataSource = DataSource.EXTRACTED_PROFILE,
+        )
+        val optOutStep = OptOutStep(
+            broker = testBroker,
+            step = OptOutStepActions(
+                stepType = "optout",
+                actions = listOf(action),
+                optOutType = "form",
+            ),
+            profileToOptOut = testExtractedProfile,
+        )
+        val state = State(
+            runType = RunType.OPTOUT,
+            brokerStepsToExecute = listOf(optOutStep),
+            profileQuery = testProfileQuery,
+            currentBrokerStepIndex = 0,
+            currentActionIndex = 0,
+            emailExtractedData = emptyMap(),
+            stageStatus = PirStageStatus(
+                currentStage = PirStage.OTHER,
+                stageStartMs = 0,
+            ),
+        )
+        val event = ExecuteBrokerStepAction(UserProfile(userProfile = testProfileQuery))
+
+        val result = testee.invoke(state, event)
+
+        val sideEffect = result.sideEffect as PushJsAction
+        val userData = sideEffect.requestParamsData as UserProfile
+        assertNull(userData.extractedProfile?.emailExtractedData)
+    }
 }