Bug 1880476 — Messaging: promote NimbusMessagingController to components.nimbus.messaging

app/src/androidTest/java/org/mozilla/fenix/ui/NimbusMessagingMessageTest.kt

@@ -33,8 +33,8 @@ class NimbusMessagingMessageTest {
 
     private lateinit var context: Context
 
-    private val storage
+    private val messaging
-        get() = context.components.nimbus.messagingStorage
+        get() = context.components.nimbus.messaging
 
     @get:Rule
     val activityTestRule =
@@ -54,7 +54,7 @@ class NimbusMessagingMessageTest {
      */
     @Test
     fun testAllMessageIntegrity() = runTest {
-        val messages = storage.getMessages()
+        val messages = messaging.getMessages()
         val rawMessages = feature.messages
         assertTrue(rawMessages.isNotEmpty())
 
@@ -67,22 +67,6 @@ class NimbusMessagingMessageTest {
         assertEquals(messages.size, rawMessages.size)
     }
 
-    /**
-     * Check if the messages' triggers are well formed JEXL.
-     */
-    @Test
-    fun testAllMessageTriggers() = runTest {
-        val helper = context.components.nimbus.createJexlHelper()
-        val messages = storage.getMessages()
-        messages.forEach { message ->
-            storage.isMessageEligible(message, helper)
-            if (storage.malFormedMap.isNotEmpty()) {
-                fail("${message.id} has a problem with its JEXL trigger: ${storage.malFormedMap.keys}")
-            }
-        }
-        helper.destroy()
-    }
-
     private fun checkIsLocalized(string: String) {
         assertFalse(string.isBlank())
         // The check will almost always succeed, since the generated code

app/src/main/java/org/mozilla/fenix/HomeActivity.kt

@@ -121,7 +121,6 @@ import org.mozilla.fenix.home.intent.SpeechProcessingIntentProcessor
 import org.mozilla.fenix.home.intent.StartSearchIntentProcessor
 import org.mozilla.fenix.library.bookmarks.DesktopFolders
 import org.mozilla.fenix.messaging.FenixMessageSurfaceId
-import org.mozilla.fenix.messaging.FenixNimbusMessagingController
 import org.mozilla.fenix.messaging.MessageNotificationWorker
 import org.mozilla.fenix.nimbus.FxNimbus
 import org.mozilla.fenix.onboarding.ReEngagementNotificationWorker
@@ -1221,13 +1220,8 @@ open class HomeActivity : LocaleAwareAppCompatActivity(), NavHostActivity {
     }
 
     private suspend fun showFullscreenMessageIfNeeded(context: Context) {
-        val messagingStorage = context.components.nimbus.messagingStorage
+        val messaging = context.components.nimbus.messaging
-        val messages = messagingStorage.getMessages()
+        val nextMessage = messaging.getNextMessage(FenixMessageSurfaceId.SURVEY) ?: return
-        val nextMessage =
-            messagingStorage.getNextMessage(FenixMessageSurfaceId.SURVEY, messages)
-                ?: return
-
-        val fenixNimbusMessagingController = FenixNimbusMessagingController(messagingStorage)
         val researchSurfaceDialogFragment = ResearchSurfaceDialogFragment.newInstance(
             keyMessageText = nextMessage.text,
             keyAcceptButtonText = nextMessage.buttonLabel,
@@ -1235,7 +1229,7 @@ open class HomeActivity : LocaleAwareAppCompatActivity(), NavHostActivity {
         )
 
         researchSurfaceDialogFragment.onAccept = {
-            processIntent(fenixNimbusMessagingController.getIntentForMessage(nextMessage))
+            processIntent(messaging.getIntentForMessage(nextMessage))
             components.appStore.dispatch(AppAction.MessagingAction.MessageClicked(nextMessage))
         }
 
@@ -1252,15 +1246,8 @@ open class HomeActivity : LocaleAwareAppCompatActivity(), NavHostActivity {
 
         // Update message as displayed.
         val currentBootUniqueIdentifier = BootUtils.getBootIdentifier(context)
-        val updatedMessage =
-            fenixNimbusMessagingController.updateMessageAsDisplayed(
-                nextMessage,
-                currentBootUniqueIdentifier,
-            )
-
-        fenixNimbusMessagingController.onMessageDisplayed(updatedMessage)
 
-        return
+        messaging.onMessageDisplayed(nextMessage, currentBootUniqueIdentifier)
     }
 
     companion object {

app/src/main/java/org/mozilla/fenix/components/Components.kt

@@ -232,7 +232,7 @@ class Components(private val context: Context) {
                     context.pocketStoriesSelectedCategoriesDataStore,
                 ),
                 MessagingMiddleware(
-                    messagingStorage = nimbus.messagingStorage,
+                    controller = nimbus.messaging,
                 ),
                 MetricsMiddleware(metrics = analytics.metrics),
             ),

app/src/main/java/org/mozilla/fenix/components/NimbusComponents.kt

@@ -7,6 +7,8 @@ package org.mozilla.fenix.components
 import android.content.Context
 import mozilla.components.service.nimbus.NimbusApi
 import mozilla.components.service.nimbus.messaging.FxNimbusMessaging
+import mozilla.components.service.nimbus.messaging.NimbusMessagingController
+import mozilla.components.service.nimbus.messaging.NimbusMessagingControllerInterface
 import mozilla.components.service.nimbus.messaging.NimbusMessagingStorage
 import mozilla.components.service.nimbus.messaging.OnDiskMessageMetadataStorage
 import org.mozilla.experiments.nimbus.NimbusEventStore
@@ -71,12 +73,23 @@ class NimbusComponents(private val context: Context) {
     fun createJexlHelper(): NimbusMessagingHelperInterface =
         messagingStorage.createMessagingHelper()
 
+    /**
+     * The main entry point for UI surfaces to interact with (get, click, dismiss) messages
+     * from the Nimbus Messaging component.
+     */
+    val messaging: NimbusMessagingControllerInterface by lazyMonitored {
+        NimbusMessagingController(
+            messagingStorage = messagingStorage,
+            deepLinkScheme = BuildConfig.DEEP_LINK_SCHEME,
+        )
+    }
+
     /**
      * Low level access to the messaging component.
      *
      * The app should access this through a [mozilla.components.service.nimbus.messaging.NimbusMessagingController].
      */
-    val messagingStorage by lazyMonitored {
+    private val messagingStorage by lazyMonitored {
         NimbusMessagingStorage(
             context = context,
             metadataStorage = OnDiskMessageMetadataStorage(context),

app/src/main/java/org/mozilla/fenix/home/HomeFragment.kt

@@ -117,7 +117,6 @@ import org.mozilla.fenix.home.toolbar.SearchSelectorBinding
 import org.mozilla.fenix.home.toolbar.SearchSelectorMenuBinding
 import org.mozilla.fenix.home.topsites.DefaultTopSitesView
 import org.mozilla.fenix.messaging.DefaultMessageController
-import org.mozilla.fenix.messaging.FenixNimbusMessagingController
 import org.mozilla.fenix.messaging.MessagingFeature
 import org.mozilla.fenix.nimbus.FxNimbus
 import org.mozilla.fenix.perf.MarkersFragmentLifecycleCallbacks
@@ -360,7 +359,7 @@ class HomeFragment : Fragment() {
                 engine = components.core.engine,
                 messageController = DefaultMessageController(
                     appStore = components.appStore,
-                    messagingController = FenixNimbusMessagingController(components.nimbus.messagingStorage),
+                    messagingController = components.nimbus.messaging,
                     homeActivity = activity,
                 ),
                 store = store,

app/src/main/java/org/mozilla/fenix/messaging/DefaultMessageController.kt

@@ -4,9 +4,8 @@
 
 package org.mozilla.fenix.messaging
 
-import android.content.Intent
 import mozilla.components.service.nimbus.messaging.Message
-import mozilla.components.service.nimbus.messaging.NimbusMessagingController
+import mozilla.components.service.nimbus.messaging.NimbusMessagingControllerInterface
 import org.mozilla.fenix.HomeActivity
 import org.mozilla.fenix.components.AppStore
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.MessageClicked
@@ -17,13 +16,13 @@ import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.MessageDi
  */
 class DefaultMessageController(
     private val appStore: AppStore,
-    private val messagingController: NimbusMessagingController,
+    private val messagingController: NimbusMessagingControllerInterface,
     private val homeActivity: HomeActivity,
 ) : MessageController {
 
     override fun onMessagePressed(message: Message) {
-        val actionUri = messagingController.processMessageActionToUri(message)
+        val intent = messagingController.getIntentForMessage(message)
-        homeActivity.processIntent(Intent(Intent.ACTION_VIEW, actionUri))
+        homeActivity.processIntent(intent)
 
         appStore.dispatch(MessageClicked(message))
     }

app/src/main/java/org/mozilla/fenix/messaging/FenixNimbusMessagingController.kt

@@ -1,23 +0,0 @@
-/* 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.messaging
-
-import mozilla.components.service.nimbus.messaging.NimbusMessagingController
-import mozilla.components.service.nimbus.messaging.NimbusMessagingStorage
-import org.mozilla.fenix.BuildConfig
-
-/**
- * Bookkeeping for message actions in terms of Glean messages and the messaging store.
- * Specialized implementation of NimbusMessagingController defining the deepLinkScheme
- * used by Fenix
- */
-class FenixNimbusMessagingController(
-    messagingStorage: NimbusMessagingStorage,
-    now: () -> Long = { System.currentTimeMillis() },
-) : NimbusMessagingController(
-    messagingStorage = messagingStorage,
-    deepLinkScheme = BuildConfig.DEEP_LINK_SCHEME,
-    now = now,
-)

app/src/main/java/org/mozilla/fenix/messaging/MessageNotificationWorker.kt

@@ -45,10 +45,10 @@ class MessageNotificationWorker(
     override suspend fun doWork(): Result {
         val context = applicationContext
 
-        val messagingStorage = context.components.nimbus.messagingStorage
+        val messaging = context.components.nimbus.messaging
-        val messages = messagingStorage.getMessages()
+
         val nextMessage =
-            messagingStorage.getNextMessage(FenixMessageSurfaceId.NOTIFICATION, messages)
+            messaging.getNextMessage(FenixMessageSurfaceId.NOTIFICATION)
                 ?: return Result.success()
 
         val currentBootUniqueIdentifier = BootUtils.getBootIdentifier(context)
@@ -57,23 +57,15 @@ class MessageNotificationWorker(
             return Result.success()
         }
 
-        val nimbusMessagingController = FenixNimbusMessagingController(messagingStorage)
-
         // Update message as displayed.
-        val updatedMessage =
+        messaging.onMessageDisplayed(nextMessage, currentBootUniqueIdentifier)
-            nimbusMessagingController.updateMessageAsDisplayed(
-                nextMessage,
-                currentBootUniqueIdentifier,
-            )
-
-        nimbusMessagingController.onMessageDisplayed(updatedMessage)
 
         context.components.notificationsDelegate.notify(
             MESSAGE_TAG,
-            SharedIdsHelper.getIdForTag(context, updatedMessage.id),
+            SharedIdsHelper.getIdForTag(context, nextMessage.id),
             buildNotification(
                 context,
-                updatedMessage,
+                nextMessage,
             ),
         )
 
@@ -177,18 +169,17 @@ class NotificationDismissedService : LifecycleService() {
         super.onStartCommand(intent, flags, startId)
 
         if (intent != null) {
-            val nimbusMessagingController =
+            val messaging = applicationContext.components.nimbus.messaging
-                FenixNimbusMessagingController(applicationContext.components.nimbus.messagingStorage)
 
             lifecycleScope.launch {
                 // Get the relevant message.
                 val message = intent.getStringExtra(DISMISSED_MESSAGE_ID)?.let { messageId ->
-                    nimbusMessagingController.getMessage(messageId)
+                    messaging.getMessage(messageId)
                 }
 
                 if (message != null) {
                     // Update message as 'dismissed'.
-                    nimbusMessagingController.onMessageDismissed(message)
+                    messaging.onMessageDismissed(message)
                 }
             }
         }
@@ -208,21 +199,20 @@ class NotificationClickedReceiverActivity : ComponentActivity() {
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
 
-        val nimbusMessagingController =
+        val messaging = applicationContext.components.nimbus.messaging
-            FenixNimbusMessagingController(components.nimbus.messagingStorage)
 
         lifecycleScope.launch {
             // Get the relevant message.
             val message = intent.getStringExtra(CLICKED_MESSAGE_ID)?.let { messageId ->
-                nimbusMessagingController.getMessage(messageId)
+                messaging.getMessage(messageId)
             }
 
             if (message != null) {
                 // Update message as 'clicked'.
-                nimbusMessagingController.onMessageClicked(message)
+                messaging.onMessageClicked(message)
 
                 // Create the intent.
-                val intent = nimbusMessagingController.getIntentForMessage(message)
+                val intent = messaging.getIntentForMessage(message)
                 intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
                 intent.addFlags(Intent.FLAG_ACTIVITY_CLEAR_TASK)
 

app/src/main/java/org/mozilla/fenix/messaging/state/MessagingMiddleware.kt

@@ -10,8 +10,7 @@ import kotlinx.coroutines.launch
 import mozilla.components.lib.state.Middleware
 import mozilla.components.lib.state.MiddlewareContext
 import mozilla.components.service.nimbus.messaging.Message
-import mozilla.components.service.nimbus.messaging.NimbusMessagingController
+import mozilla.components.service.nimbus.messaging.NimbusMessagingControllerInterface
-import mozilla.components.service.nimbus.messaging.NimbusMessagingStorage
 import org.mozilla.fenix.components.appstate.AppAction
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.ConsumeMessageToShow
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.Evaluate
@@ -21,13 +20,11 @@ import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.Restore
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.UpdateMessageToShow
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.UpdateMessages
 import org.mozilla.fenix.components.appstate.AppState
-import org.mozilla.fenix.messaging.FenixNimbusMessagingController
 
 typealias AppStoreMiddlewareContext = MiddlewareContext<AppState, AppAction>
 
 class MessagingMiddleware(
-    private val messagingStorage: NimbusMessagingStorage,
+    private val controller: NimbusMessagingControllerInterface,
-    private val controller: NimbusMessagingController = FenixNimbusMessagingController(messagingStorage),
     private val coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.IO),
 ) : Middleware<AppState, AppAction> {
 
@@ -39,13 +36,13 @@ class MessagingMiddleware(
         when (action) {
             is Restore -> {
                 coroutineScope.launch {
-                    val messages = messagingStorage.getMessages()
+                    val messages = controller.getMessages()
                     context.store.dispatch(UpdateMessages(messages))
                 }
             }
 
             is Evaluate -> {
-                val message = messagingStorage.getNextMessage(
+                val message = controller.getNextMessage(
                     action.surface,
                     context.state.messaging.messages,
                 )
@@ -72,16 +69,15 @@ class MessagingMiddleware(
         oldMessage: Message,
         context: AppStoreMiddlewareContext,
     ) {
-        val newMessage = controller.updateMessageAsDisplayed(oldMessage)
-        val newMessages = if (!newMessage.isExpired) {
-            updateMessage(context, oldMessage, newMessage)
-        } else {
-            consumeMessageToShowIfNeeded(context, oldMessage)
-            removeMessage(context, oldMessage)
-        }
-        context.dispatch(UpdateMessages(newMessages))
         coroutineScope.launch {
-            controller.onMessageDisplayed(newMessage)
+            val newMessage = controller.onMessageDisplayed(oldMessage)
+            val newMessages = if (!newMessage.isExpired) {
+                updateMessage(context, oldMessage, newMessage)
+            } else {
+                consumeMessageToShowIfNeeded(context, oldMessage)
+                removeMessage(context, oldMessage)
+            }
+            context.store.dispatch(UpdateMessages(newMessages))
         }
     }
 
@@ -136,7 +132,7 @@ class MessagingMiddleware(
         val actualMessageToShow = context.state.messaging.messageToShow[updatedMessage.surface]
 
         if (actualMessageToShow?.id == oldMessage.id) {
-            context.dispatch(UpdateMessageToShow(updatedMessage))
+            context.store.dispatch(UpdateMessageToShow(updatedMessage))
         }
         val oldMessageIndex = context.state.messaging.messages.indexOfFirst { it.id == updatedMessage.id }
         val newList = context.state.messaging.messages.toMutableList()

app/src/test/java/org/mozilla/fenix/messaging/DefaultMessageControllerTest.kt

@@ -4,13 +4,11 @@
 
 package org.mozilla.fenix.messaging
 
-import androidx.core.net.toUri
-import io.mockk.every
 import io.mockk.mockk
 import io.mockk.verify
 import mozilla.components.service.nimbus.messaging.Message
 import mozilla.components.service.nimbus.messaging.MessageData
-import mozilla.components.service.nimbus.messaging.NimbusMessagingController
+import mozilla.components.service.nimbus.messaging.NimbusMessagingControllerInterface
 import mozilla.components.support.test.robolectric.testContext
 import mozilla.telemetry.glean.testing.GleanTestRule
 import org.junit.Before
@@ -30,7 +28,7 @@ class DefaultMessageControllerTest {
     val gleanTestRule = GleanTestRule(testContext)
 
     private val homeActivity: HomeActivity = mockk(relaxed = true)
-    private val messagingController: NimbusMessagingController = mockk(relaxed = true)
+    private val messagingController: NimbusMessagingControllerInterface = mockk(relaxed = true)
     private lateinit var defaultMessageController: DefaultMessageController
     private val appStore: AppStore = mockk(relaxed = true)
 
@@ -46,12 +44,10 @@ class DefaultMessageControllerTest {
     @Test
     fun `WHEN calling onMessagePressed THEN process the action intent and update the app store`() {
         val message = mockMessage()
-        val uri = "action".toUri()
-        every { messagingController.processMessageActionToUri(message) }.returns(uri)
 
         defaultMessageController.onMessagePressed(message)
 
-        verify { messagingController.processMessageActionToUri(message) }
+        verify { messagingController.getIntentForMessage(message) }
         verify { homeActivity.processIntent(any()) }
         verify { appStore.dispatch(MessageClicked(message)) }
     }

app/src/test/java/org/mozilla/fenix/messaging/state/MessagingMiddlewareTest.kt

@@ -12,13 +12,13 @@ import kotlinx.coroutines.test.advanceUntilIdle
 import mozilla.components.service.nimbus.messaging.Message
 import mozilla.components.service.nimbus.messaging.MessageData
 import mozilla.components.service.nimbus.messaging.NimbusMessagingController
-import mozilla.components.service.nimbus.messaging.NimbusMessagingStorage
 import mozilla.components.service.nimbus.messaging.StyleData
 import mozilla.components.support.test.ext.joinBlocking
 import mozilla.components.support.test.libstate.ext.waitUntilIdle
 import mozilla.components.support.test.rule.MainCoroutineRule
 import mozilla.components.support.test.rule.runTestOnMain
 import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
 import org.junit.Before
 import org.junit.Rule
@@ -30,20 +30,17 @@ import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.MessageDi
 import org.mozilla.fenix.components.appstate.AppAction.MessagingAction.Restore
 import org.mozilla.fenix.components.appstate.AppState
 import org.mozilla.fenix.messaging.FenixMessageSurfaceId
-import org.mozilla.fenix.messaging.FenixNimbusMessagingController
 import org.mozilla.fenix.messaging.MessagingState
 
 class MessagingMiddlewareTest {
     @get:Rule
     val coroutinesTestRule = MainCoroutineRule()
     private val coroutineScope = coroutinesTestRule.scope
-    private lateinit var messagingStorage: NimbusMessagingStorage
     private lateinit var controller: NimbusMessagingController
 
     @Before
     fun setUp() {
-        messagingStorage = mockk(relaxed = true)
+        controller = mockk(relaxed = true)
-        controller = FenixNimbusMessagingController(messagingStorage) { 0 }
     }
 
     @Test
@@ -55,12 +52,12 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
         val message = createMessage()
 
-        coEvery { messagingStorage.getMessages() } returns listOf(message)
+        coEvery { controller.getMessages() } returns listOf(message)
 
         store.dispatch(Restore).joinBlocking()
         store.waitUntilIdle()
@@ -81,12 +78,12 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
         every {
-            messagingStorage.getNextMessage(
+            controller.getNextMessage(
                 FenixMessageSurfaceId.HOMESCREEN,
                 any(),
             )
@@ -111,7 +108,7 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
@@ -121,7 +118,7 @@ class MessagingMiddlewareTest {
         store.waitUntilIdle()
 
         assertTrue(store.state.messaging.messages.isEmpty())
-        coVerify { messagingStorage.updateMetadata(createMetadata(pressed = true)) }
+        coVerify { controller.onMessageClicked(message = message) }
     }
 
     @Test
@@ -135,14 +132,14 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
         store.dispatch(MessageDismissed(message)).joinBlocking()
         store.waitUntilIdle()
 
         assertTrue(store.state.messaging.messages.isEmpty())
-        coVerify { messagingStorage.updateMetadata(metadata.copy(dismissed = true)) }
+        coVerify { controller.onMessageDismissed(message = message) }
     }
 
     @Test
@@ -158,7 +155,7 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
@@ -184,7 +181,7 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
@@ -198,6 +195,7 @@ class MessagingMiddlewareTest {
     @Test
     fun `WHEN updateMessage THEN update available messages`() = runTestOnMain {
         val message = createMessage()
+        val messageDisplayed = message.copy(metadata = createMetadata(displayCount = 1))
         val store = AppStore(
             AppState(
                 messaging = MessagingState(
@@ -207,20 +205,25 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
         every {
-            messagingStorage.getNextMessage(
+            controller.getNextMessage(
                 FenixMessageSurfaceId.HOMESCREEN,
                 any(),
             )
         } returns message
 
-        store.dispatch(Evaluate(FenixMessageSurfaceId.HOMESCREEN))
+        coEvery {
+            controller.onMessageDisplayed(eq(message), any())
+        } returns messageDisplayed
+
+        store.dispatch(Evaluate(FenixMessageSurfaceId.HOMESCREEN)).joinBlocking()
         store.waitUntilIdle()
 
+        assertEquals(1, store.state.messaging.messages.count())
         assertEquals(1, store.state.messaging.messages.first().displayCount)
     }
 
@@ -240,17 +243,21 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
         every {
-            messagingStorage.getNextMessage(
+            controller.getNextMessage(
                 FenixMessageSurfaceId.HOMESCREEN,
                 any(),
             )
         } returns message1
 
+        coEvery {
+            controller.onMessageDisplayed(eq(message1), any())
+        } returns messageDisplayed1
+
         store.dispatch(Evaluate(FenixMessageSurfaceId.HOMESCREEN)).joinBlocking()
         store.waitUntilIdle()
 
@@ -272,17 +279,21 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
         every {
-            messagingStorage.getNextMessage(
+            controller.getNextMessage(
                 FenixMessageSurfaceId.HOMESCREEN,
                 any(),
             )
         } returns message
 
+        coEvery {
+            controller.onMessageDisplayed(eq(message), any())
+        } returns messageDisplayed
+
         store.dispatch(Evaluate(FenixMessageSurfaceId.HOMESCREEN)).joinBlocking()
         store.waitUntilIdle()
 
@@ -291,7 +302,11 @@ class MessagingMiddlewareTest {
 
     @Test
     fun `GIVEN a message with that surpassed the maxDisplayCount WHEN onMessagedDisplayed THEN remove the message and consume it`() = runTestOnMain {
-        val message = createMessage(createMetadata(displayCount = 6))
+        val message = createMessage(createMetadata(displayCount = 4))
+        val messageDisplayed = createMessage(createMetadata(displayCount = 5))
+        assertFalse(message.isExpired)
+        assertTrue(messageDisplayed.isExpired)
+
         val store = AppStore(
             AppState(
                 messaging = MessagingState(
@@ -302,17 +317,21 @@ class MessagingMiddlewareTest {
                 ),
             ),
             listOf(
-                MessagingMiddleware(messagingStorage, controller, coroutineScope),
+                MessagingMiddleware(controller, coroutineScope),
             ),
         )
 
         every {
-            messagingStorage.getNextMessage(
+            controller.getNextMessage(
                 FenixMessageSurfaceId.HOMESCREEN,
                 any(),
             )
         } returns message
 
+        coEvery {
+            controller.onMessageDisplayed(eq(message), any())
+        } returns messageDisplayed
+
         store.dispatch(Evaluate(FenixMessageSurfaceId.HOMESCREEN)).joinBlocking()
         store.waitUntilIdle()
 
@@ -322,7 +341,7 @@ class MessagingMiddlewareTest {
 }
 private fun createMessage(
     metadata: Message.Metadata = createMetadata(),
-    messageId: String = "control-id",
+    messageId: String = "message-id",
     data: MessageData = mockk(relaxed = true),
     action: String = "action",
     styleData: StyleData = StyleData(),