For #17117: Use proper custom tab session for custom tab toolbar menu.

app/src/main/java/org/mozilla/fenix/customtabs/CustomTabToolbarMenu.kt

@@ -7,6 +7,7 @@ package org.mozilla.fenix.customtabs
 import android.content.Context
 import android.graphics.Typeface
 import androidx.annotation.ColorRes
+import androidx.annotation.VisibleForTesting
 import androidx.core.content.ContextCompat.getColor
 import mozilla.components.browser.menu.BrowserMenuBuilder
 import mozilla.components.browser.menu.BrowserMenuHighlight
@@ -17,8 +18,8 @@ import mozilla.components.browser.menu.item.BrowserMenuImageSwitch
 import mozilla.components.browser.menu.item.BrowserMenuImageText
 import mozilla.components.browser.menu.item.BrowserMenuItemToolbar
 import mozilla.components.browser.menu.item.SimpleBrowserMenuItem
-import mozilla.components.browser.state.selector.findTab
-import mozilla.components.browser.state.state.TabSessionState
+import mozilla.components.browser.state.selector.findCustomTab
+import mozilla.components.browser.state.state.CustomTabSessionState
 import mozilla.components.browser.state.store.BrowserStore
 import org.mozilla.fenix.R
 import org.mozilla.fenix.components.toolbar.ToolbarMenu
@@ -46,7 +47,8 @@ class CustomTabToolbarMenu(
     override val menuBuilder by lazy { BrowserMenuBuilder(menuItems) }
 
     /** Gets the current custom tab session */
-    private val session: TabSessionState? get() = sessionId?.let { store.state.findTab(it) }
+    @VisibleForTesting
+    internal val session: CustomTabSessionState? get() = sessionId?.let { store.state.findCustomTab(it) }
     private val appName = context.getString(R.string.app_name)
 
     override val menuToolbar by lazy {

app/src/test/java/org/mozilla/fenix/customtabs/CustomTabToolbarMenuTest.kt

@@ -0,0 +1,62 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.fenix.customtabs
+
+import android.content.Context
+import io.mockk.mockk
+import io.mockk.spyk
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.CustomTabSessionState
+import mozilla.components.browser.state.state.createCustomTab
+import mozilla.components.browser.state.state.createTab
+import mozilla.components.browser.state.store.BrowserStore
+import org.junit.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+
+@ExperimentalCoroutinesApi
+class CustomTabToolbarMenuTest {
+
+    private lateinit var firefoxCustomTab: CustomTabSessionState
+    private lateinit var store: BrowserStore
+    private lateinit var customTabToolbarMenu: CustomTabToolbarMenu
+    private lateinit var context: Context
+
+    @Before
+    fun setUp() {
+        context = mockk(relaxed = true)
+
+        firefoxCustomTab = createCustomTab(url = "https://firefox.com", id = "123")
+
+        store = BrowserStore(
+            BrowserState(
+                tabs = listOf(
+                    createTab(url = "https://wikipedia.com", id = "1")
+                ),
+                customTabs = listOf(
+                    firefoxCustomTab,
+                    createCustomTab(url = "https://mozilla.com", id = "456")
+                )
+            )
+        )
+
+        customTabToolbarMenu = spyk(
+            CustomTabToolbarMenu(
+                context = context,
+                store = store,
+                sessionId = firefoxCustomTab.id,
+                shouldReverseItems = false,
+                onItemTapped = { }
+            )
+        )
+    }
+
+    @Test
+    fun `custom tab toolbar menu uses the proper custom tab session`() {
+        assertEquals(firefoxCustomTab.id, customTabToolbarMenu.session?.id)
+        assertEquals("https://firefox.com", customTabToolbarMenu.session?.content?.url)
+    }
+}