From 7f3ca1010308f5617344db14c4b5c2bd312b6590 Mon Sep 17 00:00:00 2001
From: ekager <ekager@mozilla.com>
Date: Wed, 15 Apr 2020 22:56:27 -0700
Subject: [PATCH] [fenix] For
 https://github.com/mozilla-mobile/fenix/issues/9981 - Make Bookmark Menu UI
 respond to session selection

---
 .../components/toolbar/DefaultToolbarMenu.kt  | 35 ++++++++++++++++---
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/app/src/main/java/org/mozilla/fenix/components/toolbar/DefaultToolbarMenu.kt b/app/src/main/java/org/mozilla/fenix/components/toolbar/DefaultToolbarMenu.kt
index 702d89379d..1bed2b308e 100644
--- a/app/src/main/java/org/mozilla/fenix/components/toolbar/DefaultToolbarMenu.kt
+++ b/app/src/main/java/org/mozilla/fenix/components/toolbar/DefaultToolbarMenu.kt
@@ -203,7 +203,11 @@ class DefaultToolbarMenu(
             menuToolbar
         )
 
-        if (shouldReverseItems) { menuItems.reversed() } else { menuItems }
+        if (shouldReverseItems) {
+            menuItems.reversed()
+        } else {
+            menuItems
+        }
     }
 
     private val addons = BrowserMenuImageText(
@@ -355,16 +359,39 @@ class DefaultToolbarMenu(
     @ColorRes
     private fun primaryTextColor() = ThemeManager.resolveAttribute(R.attr.primaryText, context)
 
+    private var currentSessionObserver: Pair<Session, Session.Observer>? = null
+
     private fun registerForIsBookmarkedUpdates() {
-        val observer = object : Session.Observer {
+        session?.let {
+            registerForUrlChanges(it)
+        }
+
+        val sessionManagerObserver = object : SessionManager.Observer {
+            override fun onSessionSelected(session: Session) {
+                // Unregister any old session observer before registering a new session observer
+                currentSessionObserver?.let {
+                    it.first.unregister(it.second)
+                }
+                currentUrlIsBookmarked = false
+                updateCurrentUrlIsBookmarked(session.url)
+                registerForUrlChanges(session)
+            }
+        }
+
+        sessionManager.register(sessionManagerObserver, lifecycleOwner)
+    }
+
+    private fun registerForUrlChanges(session: Session) {
+        val sessionObserver = object : Session.Observer {
             override fun onUrlChanged(session: Session, url: String) {
                 currentUrlIsBookmarked = false
                 updateCurrentUrlIsBookmarked(url)
             }
         }
 
-        session?.url?.let { updateCurrentUrlIsBookmarked(it) }
-        session?.register(observer, lifecycleOwner)
+        currentSessionObserver = Pair(session, sessionObserver)
+        updateCurrentUrlIsBookmarked(session.url)
+        session.register(sessionObserver, lifecycleOwner)
     }
 
     private fun updateCurrentUrlIsBookmarked(newUrl: String) {