diff --git a/components/feature/tabs/src/main/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenter.kt b/components/feature/tabs/src/main/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenter.kt
index 047d5aa149c..f10411918dc 100644
--- a/components/feature/tabs/src/main/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenter.kt
+++ b/components/feature/tabs/src/main/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenter.kt
@@ -31,24 +31,20 @@ class TabsTrayPresenter(
     private val closeTabsTray: () -> Unit
 ) {
     private var tabs: Tabs? = null
-    private var tabScope: CoroutineScope? = null
+    private var scope: CoroutineScope? = null
 
     fun start() {
-        tabScope = store.flowScoped { flow -> collect(flow) }
+        scope = store.flowScoped { flow -> collect(flow) }
     }
 
     fun stop() {
-        tabScope?.cancel()
+        scope?.cancel()
     }
 
     private suspend fun collect(flow: Flow<BrowserState>) {
-        flow.map { state ->
-            val tabs = state.toTabs(tabsFilter)
-            tabs
-        }
+        flow.map { it.toTabs(tabsFilter) }
             .ifChanged()
             .collect { tabs ->
-
                 // Do not invoke the callback on start if this is the initial state.
                 if (tabs.list.isEmpty() && this.tabs != null) {
                     closeTabsTray.invoke()
diff --git a/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/ext/TabSessionStateTest.kt b/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/ext/TabSessionStateTest.kt
new file mode 100644
index 00000000000..4fe059b37f5
--- /dev/null
+++ b/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/ext/TabSessionStateTest.kt
@@ -0,0 +1,40 @@
+/* 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 mozilla.components.feature.tabs.ext
+
+import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.MediaState
+import mozilla.components.browser.state.state.createTab
+import mozilla.components.concept.engine.media.Media
+import org.junit.Test
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNull
+
+class TabSessionStateTest {
+    @Test
+    fun `mediaState gets set correctly`() {
+        var browserState = BrowserState(
+            tabs = listOf(
+                createTab("https://www.mozilla.org", id = "a"),
+                createTab("https://getpocket.com", id = "b"),
+                createTab("https://developer.mozilla.org", id = "c"),
+                createTab("https://www.firefox.com", id = "d"),
+                createTab("https://www.google.com", id = "e")
+            ),
+            selectedTabId = "a",
+            media = MediaState(MediaState.Aggregate(activeTabId = "a", state = MediaState.State.PLAYING))
+        )
+
+        var tabs = browserState.toTabs()
+        assertEquals(Media.State.PLAYING, tabs.list[0].mediaState)
+
+        browserState = browserState.copy(
+            media = MediaState(MediaState.Aggregate(activeTabId = "b", state = MediaState.State.PAUSED))
+        )
+        tabs = browserState.toTabs()
+        assertNull(tabs.list[0].mediaState)
+        assertEquals(Media.State.PAUSED, tabs.list[1].mediaState)
+    }
+}
diff --git a/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenterTest.kt b/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenterTest.kt
index 7297a197fba..3af3e74436b 100644
--- a/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenterTest.kt
+++ b/components/feature/tabs/src/test/java/mozilla/components/feature/tabs/tabstray/TabsTrayPresenterTest.kt
@@ -222,15 +222,15 @@ class TabsTrayPresenterTest {
     @Test
     fun `tabs tray will get updated if mediaState changes`() {
         val store = BrowserStore(
-                BrowserState(
-                        tabs = listOf(
-                                createTab("https://www.mozilla.org", id = "a"),
-                                createTab("https://getpocket.com", id = "b"),
-                                createTab("https://developer.mozilla.org", id = "c"),
-                                createTab("https://www.firefox.com", id = "d"),
-                                createTab("https://www.google.com", id = "e")
-                        ),
-                        selectedTabId = "a"
+            BrowserState(
+                tabs = listOf(
+                        createTab("https://www.mozilla.org", id = "a"),
+                        createTab("https://getpocket.com", id = "b"),
+                        createTab("https://developer.mozilla.org", id = "c"),
+                        createTab("https://www.firefox.com", id = "d"),
+                        createTab("https://www.google.com", id = "e")
+                    ),
+                    selectedTabId = "a"
                 )
         )
 
@@ -242,7 +242,7 @@ class TabsTrayPresenterTest {
 
         store.dispatch(
                 MediaAction.UpdateMediaAggregateAction(
-                        store.state.media.aggregate.copy(activeTabId = "a", state = MediaState.State.PLAYING)
+                    store.state.media.aggregate.copy(activeTabId = "a", state = MediaState.State.PLAYING)
                 )
         ).joinBlocking()