Detect silently dropped WebSocket connections and retry until subscriptions are restored

app/src/main/kotlin/io/homeassistant/companion/android/websocket/WebsocketManager.kt

@@ -38,6 +38,7 @@ import io.homeassistant.companion.android.settings.SettingsActivity
 import io.homeassistant.companion.android.util.hasActiveConnection
 import java.util.concurrent.TimeUnit
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
@@ -107,6 +108,9 @@ class WebsocketManager(appContext: Context, workerParams: WorkerParameters) :
     private val checkLocalNetworkPermission: CheckLocalNetworkPermissionUseCase =
         entryPoint.checkLocalNetworkPermission()
 
+    @VisibleForTesting
+    internal var dispatcher: CoroutineDispatcher = Dispatchers.IO
+
     @EntryPoint
     @InstallIn(SingletonComponent::class)
     interface WebsocketManagerEntryPoint {
@@ -116,7 +120,7 @@ class WebsocketManager(appContext: Context, workerParams: WorkerParameters) :
         fun checkLocalNetworkPermission(): CheckLocalNetworkPermissionUseCase
     }
 
-    override suspend fun doWork(): Result = withContext(Dispatchers.IO) {
+    override suspend fun doWork(): Result = withContext(dispatcher) {
         if (!checkLocalNetworkPermission()) {
             Timber.d("Skipping websocket work: ACCESS_LOCAL_NETWORK permission missing")
             return@withContext Result.success()
@@ -137,7 +141,7 @@ class WebsocketManager(appContext: Context, workerParams: WorkerParameters) :
         // play ping pong to ensure we have a connection and server changes are handled.
         do {
             delay(30000)
-        } while (jobs.values.any { it.isActive } && isActive && shouldWeRun() && manageServerJobs(jobs, this))
+        } while (isActive && shouldWeRun() && manageServerJobs(jobs, this))
 
         jobs.forEach { it.value.cancel() }
         jobs.clear()
@@ -176,9 +180,9 @@ class WebsocketManager(appContext: Context, workerParams: WorkerParameters) :
     private suspend fun manageServerJobs(jobs: MutableMap<Int, Job>, coroutineScope: CoroutineScope): Boolean {
         val servers = serverManager.servers()
 
-        // Clean up...
-        jobs.filter { (serverId, _) ->
-            servers.none { it.id == serverId } || !shouldRunForServer(serverId)
+        // Clean up, including stopped jobs so they are started again below...
+        jobs.filter { (serverId, job) ->
+            servers.none { it.id == serverId } || !job.isActive || !shouldRunForServer(serverId)
         }
             .forEach { (serverId, job) ->
                 job.cancel()

app/src/test/kotlin/io/homeassistant/companion/android/websocket/WebsocketManagerTest.kt

@@ -10,6 +10,7 @@ import androidx.work.testing.TestListenableWorkerBuilder
 import dagger.hilt.android.EntryPointAccessors
 import dagger.hilt.android.testing.HiltTestApplication
 import io.homeassistant.companion.android.common.data.servers.ServerManager
+import io.homeassistant.companion.android.common.data.websocket.WebSocketRepository
 import io.homeassistant.companion.android.common.util.CheckLocalNetworkPermissionUseCase
 import io.homeassistant.companion.android.database.server.Server
 import io.homeassistant.companion.android.database.settings.SensorUpdateFrequencySetting
@@ -19,16 +20,23 @@ import io.homeassistant.companion.android.database.settings.WebsocketSetting
 import io.homeassistant.companion.android.notifications.MessagingManager
 import io.homeassistant.companion.android.util.hasActiveConnection
 import io.mockk.coEvery
+import io.mockk.coJustRun
 import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.mockkConstructor
 import io.mockk.mockkStatic
 import io.mockk.spyk
+import io.mockk.unmockkAll
 import io.mockk.verify
 import junit.framework.TestCase.assertEquals
 import kotlinx.coroutines.CancellationException
+import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.runTest
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -81,6 +89,11 @@ class WebsocketManagerTest {
             entryPoint
     }
 
+    @After
+    fun tearDown() {
+        unmockkAll()
+    }
+
     @Test
     fun `Given setting NEVER when job runs then finishes without checks`() = runTest {
         mockSetting(WebsocketSetting.NEVER)
@@ -213,4 +226,36 @@ class WebsocketManagerTest {
         verify(exactly = 0) { context.getSystemService(PowerManager::class.java) }
         coVerify(exactly = 0) { entryPoint.serverManager.connectionStateProvider(any()) }
     }
+
+    @Test
+    fun `Given a collector that stopped when listening for notifications then it is restarted on the next interval`() = runTest {
+        mockSetting(WebsocketSetting.ALWAYS)
+        every { context.hasActiveConnection() } returns true
+        coEvery { entryPoint.serverManager.isRegistered() } returns true
+
+        mockkConstructor(NotificationCompat.Builder::class)
+        every { anyConstructed<NotificationCompat.Builder>().build() } returns mockk<Notification>()
+
+        val repository = mockk<WebSocketRepository>(relaxed = true)
+        coEvery { entryPoint.serverManager.webSocketRepository(any()) } returns repository
+        // The first collection completes immediately (the initial subscribe failed, e.g. the
+        // server was still starting), the second stays live
+        coEvery { repository.getNotifications() } returnsMany
+            listOf(null, MutableSharedFlow<Map<String, Any>>())
+
+        val worker = spyk(TestListenableWorkerBuilder<WebsocketManager>(context).build())
+        coJustRun { worker.setForeground(any()) }
+        worker.dispatcher = StandardTestDispatcher(testScheduler)
+
+        val work = launch { worker.doWork() }
+        testScheduler.runCurrent()
+        coVerify(exactly = 1) { repository.getNotifications() }
+
+        // The next supervision interval restarts the stopped collector
+        testScheduler.advanceTimeBy(31_000)
+        testScheduler.runCurrent()
+        coVerify(exactly = 2) { repository.getNotifications() }
+
+        work.cancelAndJoin()
+    }
 }

common/src/main/kotlin/io/homeassistant/companion/android/common/data/websocket/WebSocketCore.kt

@@ -51,6 +51,17 @@ internal interface WebSocketCore {
      */
     suspend fun sendBytes(data: ByteArray): Boolean
 
+    /**
+     * Sends an application-level ping and verifies that the connection is alive.
+     *
+     * If no pong is received while the connection is believed to be established, this indicates a
+     * silently dropped connection. The function will properly handle this as a closed connection
+     * and try restoring any active subscriptions.
+     *
+     * @return `true` if a pong was received, `false` otherwise.
+     */
+    suspend fun ping(): Boolean
+
     /**
      * Start a subscription for events on the websocket connection and get a Flow for listening to
      * new messages. When there are no more listeners, the subscription will automatically be cancelled