apply draft text when canceling Input Screen action while on SERP (#6527)

Task/Issue URL:
https://app.asana.com/1/137249556945/project/1208671518894266/task/1210964284781033?focus=true

### Description
Ensures that Input Screen behavior when it comes to canceling draft text
while on SERP is consistent with production.

### Steps to test this PR
 - [x] Search for anything.
 - [x] On SERP, open the Input Screen.
 - [x] Type a new text into the input field and back out.
- [x] Verify that the drafted text replaced the previous search query in
the omnibar.

Author: LukasPaczos

app/src/main/java/com/duckduckgo/app/browser/BrowserTabFragment.kt

@@ -921,7 +921,7 @@ class BrowserTabFragment :
 
             RESULT_CANCELED -> {
                 data.getStringExtra(InputScreenActivityResultParams.CANCELED_DRAFT_PARAM)?.let { query ->
-                    omnibar.setDraftTextIfNtp(query)
+                    omnibar.setDraftTextIfNtpOrSerp(query)
                 }
             }
         }

app/src/main/java/com/duckduckgo/app/browser/omnibar/Omnibar.kt

@@ -463,7 +463,7 @@ class Omnibar(
         newOmnibar.decorate(Decoration.NewTabScrollingState(canScrollUp, canScrollDown, topOfPage))
     }
 
-    fun setDraftTextIfNtp(query: String) {
-        newOmnibar.setDraftTextIfNtp(query)
+    fun setDraftTextIfNtpOrSerp(query: String) {
+        newOmnibar.setDraftTextIfNtpOrSerp(query)
     }
 }

app/src/main/java/com/duckduckgo/app/browser/omnibar/OmnibarLayout.kt

@@ -1038,8 +1038,8 @@ open class OmnibarLayout @JvmOverloads constructor(
         omnibarTextListener?.onTrackersCountFinished()
     }
 
-    fun setDraftTextIfNtp(query: String) {
-        viewModel.setDraftTextIfNtp(query)
+    fun setDraftTextIfNtpOrSerp(query: String) {
+        viewModel.setDraftTextIfNtpOrSerp(query)
     }
 
     private fun enableTextInputClickCatcher(enabled: Boolean) {

app/src/main/java/com/duckduckgo/app/browser/omnibar/OmnibarLayoutViewModel.kt

@@ -841,8 +841,10 @@ class OmnibarLayoutViewModel @Inject constructor(
         // }
     }
 
-    fun setDraftTextIfNtp(query: String) {
-        if (_viewState.value.viewMode is NewTab) {
+    fun setDraftTextIfNtpOrSerp(query: String) {
+        val isNtp = _viewState.value.viewMode is NewTab
+        val isSerp = _viewState.value.viewMode is Browser && duckDuckGoUrlDetector.isDuckDuckGoQueryUrl(_viewState.value.url)
+        if (isNtp || isSerp) {
             _viewState.update {
                 it.copy(
                     omnibarText = query,

app/src/test/java/com/duckduckgo/app/browser/omnibar/OmnibarLayoutViewModelTest.kt

@@ -1628,6 +1628,61 @@ class OmnibarLayoutViewModelTest {
         }
     }
 
+    @Test
+    fun `when set draft and current state is NTP, then draft text applied`() = runTest {
+        testee.onViewModeChanged(ViewMode.NewTab)
+        val expected = "test"
+        testee.setDraftTextIfNtpOrSerp(expected)
+
+        testee.viewState.test {
+            val viewState = awaitItem()
+            assertEquals(expected, viewState.omnibarText)
+            assertTrue(viewState.updateOmnibarText)
+        }
+    }
+
+    @Test
+    fun `when set draft and current state is SERP, then draft text applied`() = runTest {
+        givenSiteLoaded(SERP_URL)
+        val expected = "test"
+        testee.setDraftTextIfNtpOrSerp(expected)
+
+        testee.viewState.test {
+            val viewState = awaitItem()
+            assertEquals(expected, viewState.omnibarText)
+            assertTrue(viewState.updateOmnibarText)
+        }
+    }
+
+    @Test
+    fun `when set draft and current state is a web page, then draft text not applied`() = runTest {
+        val omnibarState = OmnibarViewState(
+            navigationChange = false,
+            omnibarText = RANDOM_URL,
+            forceExpand = false,
+        )
+        testee.onExternalStateChange(StateChange.OmnibarStateChange(omnibarState))
+        testee.onExternalStateChange(
+            StateChange.LoadingStateChange(
+                LoadingViewState(
+                    isLoading = true,
+                    trackersAnimationEnabled = true,
+                    progress = 100,
+                    url = RANDOM_URL,
+                ),
+            ),
+        )
+
+        val expected = RANDOM_URL
+        testee.setDraftTextIfNtpOrSerp("some draft text")
+
+        testee.viewState.test {
+            val viewState = awaitItem()
+            assertEquals(expected, viewState.omnibarText)
+            assertTrue(viewState.updateOmnibarText)
+        }
+    }
+
     private fun givenSiteLoaded(loadedUrl: String) {
         testee.onViewModeChanged(ViewMode.Browser(loadedUrl))
         testee.onExternalStateChange(