refactor: refactor navcontroller tracking plugin teardown

android/src/main/kotlin/com/rudderstack/android/sdk/plugins/screenrecording/NavControllerActivityObserver.kt

@@ -5,22 +5,24 @@ import androidx.lifecycle.DefaultLifecycleObserver
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import com.rudderstack.android.sdk.state.NavContext
+import com.rudderstack.android.sdk.utils.runOnMainThread
+import kotlinx.coroutines.DelicateCoroutinesApi
 import java.util.concurrent.atomic.AtomicBoolean
+import com.rudderstack.android.sdk.Analytics as AndroidAnalytics
 
 /*
 * This class is used to attach an observer to a navController's activity, so that destination changes can be tracked
 * when navigating back from a previous activity or when the app is foregrounded.
 * */
+@OptIn(DelicateCoroutinesApi::class)
 internal class NavControllerActivityObserver(
     private val plugin: NavControllerTrackingPlugin,
     private val navContext: NavContext,
 ) : DefaultLifecycleObserver {
 
     private val isActivityGettingCreated = AtomicBoolean(true)
 
-    init {
-        activityLifecycle()?.addObserver(this)
-    }
+    fun find(navContext: NavContext) = navContext == this.navContext
 
     override fun onStart(owner: LifecycleOwner) {
         if (!isActivityGettingCreated.getAndSet(false)) {
@@ -37,10 +39,21 @@ internal class NavControllerActivityObserver(
     }
 
     override fun onDestroy(owner: LifecycleOwner) {
-        activityLifecycle()?.removeObserver(this)
         plugin.navContextState.dispatch(NavContext.RemoveNavContextAction(navContext))
     }
 
+    internal fun removeObserver() {
+        (plugin.analytics as? AndroidAnalytics)?.runOnMainThread {
+            activityLifecycle()?.removeObserver(this)
+        }
+    }
+
+    internal fun addObserver() {
+        (plugin.analytics as? AndroidAnalytics)?.runOnMainThread {
+            activityLifecycle()?.addObserver(this)
+        }
+    }
+
     @VisibleForTesting
     internal fun activityLifecycle(): Lifecycle? {
         return (navContext.callingActivity as? LifecycleOwner)?.lifecycle

android/src/main/kotlin/com/rudderstack/android/sdk/plugins/screenrecording/NavControllerTrackingPlugin.kt

@@ -1,16 +1,17 @@
 package com.rudderstack.android.sdk.plugins.screenrecording
 
 import android.os.Bundle
-import androidx.annotation.VisibleForTesting
 import androidx.navigation.NavController
 import androidx.navigation.NavDestination
 import com.rudderstack.android.sdk.state.NavContext
 import com.rudderstack.android.sdk.utils.automaticProperty
 import com.rudderstack.kotlin.sdk.Analytics
 import com.rudderstack.kotlin.sdk.internals.plugins.Plugin
 import com.rudderstack.kotlin.sdk.internals.statemanagement.FlowState
+import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.withContext
 import com.rudderstack.android.sdk.Configuration as AndroidConfiguration
 
 internal const val FRAGMENT_NAVIGATOR_NAME = "fragment"
@@ -25,14 +26,15 @@ internal class NavControllerTrackingPlugin(
 
     override lateinit var analytics: Analytics
 
-    @VisibleForTesting
-    internal val currentNavContexts: MutableSet<NavContext> = mutableSetOf()
+    private val currentNavContexts: MutableSet<NavContext> = mutableSetOf()
+
+    private val activityObservers: MutableList<NavControllerActivityObserver> = mutableListOf()
 
     override fun setup(analytics: Analytics) {
         super.setup(analytics)
         (analytics.configuration as? AndroidConfiguration)?.let {
             navContextState.onEach { currentState ->
-                updateNavContexts(currentState)
+                withContext(NonCancellable) { updateNavContexts(currentState) }
             }.launchIn(analytics.analyticsScope)
         }
     }
@@ -58,7 +60,7 @@ internal class NavControllerTrackingPlugin(
     private fun removeDeletedNavContexts(updatedNavContexts: Set<NavContext>) {
         val deletedNavContexts = currentNavContexts.minus(updatedNavContexts)
         deletedNavContexts.forEach { navContext ->
-            removeContext(navContext)
+            removeContextAndObserver(navContext)
         }
     }
 
@@ -69,9 +71,13 @@ internal class NavControllerTrackingPlugin(
         }
     }
 
-    private fun removeContext(navContext: NavContext) {
+    private fun removeContextAndObserver(navContext: NavContext) {
         navContext.navController.removeOnDestinationChangedListener(this)
         currentNavContexts.remove(navContext)
+
+        val observerToBeRemoved = activityObservers.firstOrNull { it.find(navContext) }
+        observerToBeRemoved?.removeObserver()
+        observerToBeRemoved?.let { activityObservers.remove(it) }
     }
 
     private fun addContextAndObserver(navContext: NavContext) {
@@ -80,10 +86,12 @@ internal class NavControllerTrackingPlugin(
         currentNavContexts.add(navContext)
 
         // adding activity observer
-        NavControllerActivityObserver(
+        val observerToBeAdded = provideNavControllerActivityObserver(
             plugin = this,
             navContext = navContext
         )
+        observerToBeAdded.addObserver()
+        activityObservers.add(observerToBeAdded)
     }
 
     private fun trackFragmentScreen(destination: NavDestination) {
@@ -107,3 +115,10 @@ internal class NavControllerTrackingPlugin(
         analytics.screen(screenName = screenName, properties = automaticProperty())
     }
 }
+
+internal fun provideNavControllerActivityObserver(
+    plugin: NavControllerTrackingPlugin,
+    navContext: NavContext
+): NavControllerActivityObserver {
+    return NavControllerActivityObserver(plugin, navContext)
+}

android/src/test/kotlin/com/rudderstack/android/sdk/plugins/screenrecording/NavControllerActivityObserverTest.kt

@@ -2,7 +2,9 @@ package com.rudderstack.android.sdk.plugins.screenrecording
 
 import androidx.activity.ComponentActivity
 import androidx.navigation.NavDestination
+import com.rudderstack.kotlin.sdk.Analytics
 import com.rudderstack.android.sdk.state.NavContext
+import com.rudderstack.android.sdk.utils.mockAnalytics
 import com.rudderstack.kotlin.sdk.internals.statemanagement.FlowState
 import io.mockk.MockKAnnotations
 import io.mockk.Runs
@@ -12,12 +14,24 @@ import io.mockk.just
 import io.mockk.mockk
 import io.mockk.unmockkAll
 import io.mockk.verify
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class NavControllerActivityObserverTest {
 
+    private val testDispatcher = StandardTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
     @MockK
     private lateinit var mockPlugin: NavControllerTrackingPlugin
 
@@ -30,23 +44,30 @@ class NavControllerActivityObserverTest {
     @MockK
     private lateinit var mockNavContextState: FlowState<Set<NavContext>>
 
+    private lateinit var mockAnalytics: Analytics
+
     private lateinit var observer: NavControllerActivityObserver
 
     @Before
     fun setup() {
         MockKAnnotations.init(this, relaxed = true)
+        Dispatchers.setMain(testDispatcher)
+
+        mockAnalytics = mockAnalytics(testScope, testDispatcher)
 
         every { mockNavContext.callingActivity } returns mockActivity
         every { mockActivity.lifecycle } returns mockk(relaxed = true)
         every { mockActivity.lifecycle.addObserver(any()) } just Runs
         every { mockPlugin.navContextState } returns mockNavContextState
+        every { mockPlugin.analytics } returns mockAnalytics
         every { mockNavContextState.dispatch(any()) } just Runs
 
         observer = NavControllerActivityObserver(mockPlugin, mockNavContext)
     }
 
     @After
     fun tearDown() {
+        Dispatchers.resetMain()
         unmockkAll()
     }
 
@@ -61,10 +82,23 @@ class NavControllerActivityObserverTest {
     }
 
     @Test
-    fun `given observer is just created, then observer is also added as LifecycleObserver`() {
+    fun `given an observer, when addObserver called, then observer is added to activity lifecycle`() = runTest {
+        observer.addObserver()
+
+        advanceUntilIdle()
+
         verify(exactly = 1) { observer.activityLifecycle()?.addObserver(observer) }
     }
 
+    @Test
+    fun `given an observer, when removeObserver called, then observer is removed from activity lifecycle`() = runTest {
+        observer.removeObserver()
+
+        advanceUntilIdle()
+
+        verify(exactly = 1) { observer.activityLifecycle()?.removeObserver(observer) }
+    }
+
     @Test
     fun `given observer already created, when onStart called again, then OnDestinationChanged is not called again`() {
         val mockDestination = mockk<NavDestination>(relaxed = true)
@@ -77,11 +111,10 @@ class NavControllerActivityObserverTest {
     }
 
     @Test
-    fun `given observer, when onDestroyed called, then observer is removed and RemoveNavContextAction dispatched`() {
+    fun `given observer, when onDestroyed called, then RemoveNavContextAction dispatched`() {
         observer.onDestroy(mockActivity)
 
         verify {
-            observer.activityLifecycle()?.removeObserver(observer)
             mockNavContextState.dispatch(
                 match { it is NavContext.RemoveNavContextAction && it.navContext == mockNavContext }
             )

android/src/test/kotlin/com/rudderstack/android/sdk/plugins/screenrecording/NavControllerTrackingPluginTest.kt

@@ -1,7 +1,6 @@
 package com.rudderstack.android.sdk.plugins.screenrecording
 
 import android.os.Bundle
-import androidx.lifecycle.LifecycleOwner
 import androidx.navigation.NavController
 import androidx.navigation.NavDestination
 import com.rudderstack.android.sdk.state.NavContext
@@ -12,11 +11,10 @@ import com.rudderstack.kotlin.sdk.internals.statemanagement.FlowState
 import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.mockk
+import io.mockk.mockkStatic
 import io.mockk.spyk
 import io.mockk.unmockkAll
 import io.mockk.verify
-import junit.framework.TestCase.assertEquals
-import junit.framework.TestCase.assertTrue
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -76,30 +74,34 @@ class NavControllerTrackingPluginTest {
     }
 
     @Test
-    fun `when navContexts are added, then addOnDestinationChangedListener called and currentNavContexts updated`() =
+    fun `when navContexts are added, then addOnDestinationChangedListener and addObserver called for navContexts`() =
         runTest {
-            val navContext1: NavContext = mockNavContext()
-            val navContext2: NavContext = mockNavContext()
+            val navContext1 = mockNavContext()
+            val navContext2 = mockNavContext()
+            mockkStatic(::provideNavControllerActivityObserver)
+            val observer1 = mockActivityObserverProvider(navContext1)
+            val observer2 = mockActivityObserverProvider(navContext2)
 
             plugin.setup(mockAnalytics)
             mockNavContextState.dispatch(NavContext.AddNavContextAction(navContext1))
             mockNavContextState.dispatch(NavContext.AddNavContextAction(navContext2))
             advanceUntilIdle()
-            assertEquals(2, plugin.currentNavContexts.size)
-            assertTrue(plugin.currentNavContexts.contains(navContext1))
-            assertTrue(plugin.currentNavContexts.contains(navContext2))
+
             verify(exactly = 1) { navContext1.navController.addOnDestinationChangedListener(plugin) }
             verify(exactly = 1) { navContext2.navController.addOnDestinationChangedListener(plugin) }
+
+            verify(exactly = 1) { observer1.addObserver() }
+            verify(exactly = 1) { observer2.addObserver() }
         }
 
     @Test
-    fun `when navControllers are removed, then removeOnDestinationChangedListener called and currentNavContexts updated`() =
+    fun `when navControllers are added and one of them is removed, then removeOnDestinationChangedListener and removeObserver called for that navContext`() =
         runTest {
             val navContext1 = mockNavContext()
             val navContext2 = mockNavContext()
-
-            every { (navContext2.callingActivity as LifecycleOwner).lifecycle } returns mockk(relaxed = true)
-            every { (navContext2.callingActivity as LifecycleOwner).lifecycle.removeObserver(any()) } returns Unit
+            mockkStatic(::provideNavControllerActivityObserver)
+            val observer1 = mockActivityObserverProvider(navContext1)
+            val observer2 = mockActivityObserverProvider(navContext2)
 
             plugin.setup(mockAnalytics)
             mockNavContextState.dispatch(NavContext.AddNavContextAction(navContext1))
@@ -108,11 +110,11 @@ class NavControllerTrackingPluginTest {
             mockNavContextState.dispatch(NavContext.RemoveNavContextAction(navContext2))
             advanceUntilIdle()
 
-            // Verify that removeOnDestinationChangedListener was called for the removed navContext2 navController
-            assertEquals(1, plugin.currentNavContexts.size)
-            assertTrue(plugin.currentNavContexts.contains(navContext1))
             verify(exactly = 1) { navContext2.navController.removeOnDestinationChangedListener(plugin) }
             verify(exactly = 0) { navContext1.navController.removeOnDestinationChangedListener(plugin) }
+
+            verify(exactly = 1) { observer2.removeObserver() }
+            verify(exactly = 0) { observer1.removeObserver() }
         }
 
     @Test
@@ -164,4 +166,16 @@ class NavControllerTrackingPluginTest {
 
         verify { mockAnalytics.screen(testRouteWithoutArgs, properties = automaticProperty()) }
     }
+
+    private fun mockActivityObserverProvider(navContext: NavContext): NavControllerActivityObserver {
+        val mockObserver: NavControllerActivityObserver = mockk(relaxed = true)
+        every {
+            provideNavControllerActivityObserver(
+                plugin,
+                navContext
+            )
+        } returns mockObserver
+        every { mockObserver.find(navContext) } returns true
+        return mockObserver
+    }
 }