For #6140 - added unit tests for the dismiss action setter

components/browser/toolbar/src/test/java/mozilla/components/browser/toolbar/display/DisplayToolbarTest.kt

@@ -685,4 +685,17 @@ class DisplayToolbarTest {
 
         assertTrue(wasDismissed)
     }
+
+    @Test
+    fun `set a dismiss lambda on the menu button`() {
+        var wasDismissed = false
+        val (_, displayToolbar) = createDisplayToolbar()
+        displayToolbar.setMenuDismissAction { wasDismissed = true }
+        val menuView = displayToolbar.views.menu
+        menuView.menuBuilder = BrowserMenuBuilder(emptyList())
+        menuView.impl.performClick()
+
+        menuView.dismissMenu()
+        assertTrue(wasDismissed)
+    }
 }

components/browser/toolbar/src/test/java/mozilla/components/browser/toolbar/display/MenuButtonTest.kt

@@ -15,11 +15,11 @@ import mozilla.components.browser.menu.item.BrowserMenuHighlightableItem
 import mozilla.components.support.test.any
 import mozilla.components.support.test.mock
 import mozilla.components.support.test.robolectric.testContext
+import mozilla.components.support.test.whenever
 import org.junit.Assert.assertEquals
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mockito.`when`
 import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.never
 import org.mockito.Mockito.spy
@@ -54,6 +54,14 @@ class MenuButtonTest {
         verify(menuButtonInternal, times(2)).visibility = View.GONE
     }
 
+    @Test
+    fun `menu button sets onDismiss action`() {
+        val action = {}
+        menuButton.setMenuDismissAction(action)
+
+        verify(menuButtonInternal).onDismiss = action
+    }
+
     @Test
     fun `icon displays dot if low highlighted item is present in menu`() {
         verify(menuButtonInternal, never()).invalidateBrowserMenu()
@@ -69,7 +77,7 @@ class MenuButtonTest {
                 isHighlighted = { isHighlighted }
             )
         )))
-        `when`(highlightMenuBuilder.build(testContext)).thenReturn(menu)
+        whenever(highlightMenuBuilder.build(testContext)).thenReturn(menu)
 
         menuButton.menuBuilder = highlightMenuBuilder
         doReturn(highlightMenuBuilder).`when`(menuButtonInternal).menuBuilder