diff --git a/app/src/main/java/org/mozilla/fenix/browser/BrowserFragment.kt b/app/src/main/java/org/mozilla/fenix/browser/BrowserFragment.kt
index a9ca81bca..441fb6c7b 100644
--- a/app/src/main/java/org/mozilla/fenix/browser/BrowserFragment.kt
+++ b/app/src/main/java/org/mozilla/fenix/browser/BrowserFragment.kt
@@ -12,7 +12,6 @@ import android.view.View
 import android.view.ViewGroup
 import android.widget.Button
 import android.widget.RadioButton
-import androidx.appcompat.widget.AppCompatImageView
 import androidx.core.content.ContextCompat
 import androidx.lifecycle.Observer
 import androidx.transition.TransitionInflater
@@ -121,12 +120,10 @@ class BrowserFragment : BaseBrowserFragment(), BackHandler {
 
         val toolbarSessionObserver = TrackingProtectionOverlay(
             context = requireContext(),
-            settings = requireContext().settings(),
-            toolbar = browserToolbarView.view,
-            trackingProtectionIcon = browserToolbarView
-                .view
-                .findViewById<AppCompatImageView>(R.id.mozac_browser_toolbar_tracking_protection_indicator)
-        )
+            settings = requireContext().settings()
+        ) {
+            browserToolbarView.view
+        }
         getSessionById()?.register(toolbarSessionObserver, this, autoPause = true)
     }
 
diff --git a/app/src/main/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlay.kt b/app/src/main/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlay.kt
index 443687cb9..ad27d8d49 100644
--- a/app/src/main/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlay.kt
+++ b/app/src/main/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlay.kt
@@ -27,8 +27,7 @@ import org.mozilla.fenix.utils.Settings
 class TrackingProtectionOverlay(
     private val context: Context,
     private val settings: Settings,
-    private val toolbar: View,
-    private val trackingProtectionIcon: View
+    private val getToolbar: () -> View
 ) : Session.Observer {
 
     override fun onLoadingStateChanged(session: Session, loading: Boolean) {
@@ -74,6 +73,10 @@ class TrackingProtectionOverlay(
         val containerHeight = layout.measuredHeight
         val triangleHeight = context.getDimenInDip(R.dimen.tp_onboarding_triangle_height).toInt()
 
+        val toolbar = getToolbar()
+        val trackingProtectionIcon: View =
+            toolbar.findViewById(R.id.mozac_browser_toolbar_tracking_protection_indicator)
+
         val xOffset = context.resources.getDimensionPixelSize(R.dimen.tp_onboarding_x_offset)
 
         // Positioning the popup above the tp anchor.
diff --git a/app/src/test/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlayTest.kt b/app/src/test/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlayTest.kt
index 446a6d642..4b230bf29 100644
--- a/app/src/test/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlayTest.kt
+++ b/app/src/test/java/org/mozilla/fenix/trackingprotection/TrackingProtectionOverlayTest.kt
@@ -15,6 +15,7 @@ import mozilla.components.support.test.robolectric.testContext
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mozilla.fenix.R
 import org.mozilla.fenix.TestApplication
 import org.mozilla.fenix.utils.Settings
 import org.robolectric.RobolectricTestRunner
@@ -39,7 +40,8 @@ class TrackingProtectionOverlayTest {
         icon = mockk(relaxed = true)
         session = mockk(relaxed = true)
 
-        overlay = TrackingProtectionOverlay(context, settings, toolbar, icon)
+        overlay = TrackingProtectionOverlay(context, settings) { toolbar }
+        every { toolbar.findViewById<View>(R.id.mozac_browser_toolbar_tracking_protection_indicator) } returns icon
     }
 
     @Test