chore: improve test coverage for ConnectionStateViewModel

packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/fixtures/IntegrationTestExtensions.kt

@@ -441,7 +441,11 @@ internal fun mockLogMessage() =
  * @param url
  * @return
  */
-internal fun mockDataSource(url: MongoDbServerUrl = MongoDbServerUrl("mongodb://localhost:27017")) =
+internal fun mockDataSource(
+    url: MongoDbServerUrl = MongoDbServerUrl("mongodb://localhost:27017"),
+    name: String? = null,
+    uniqueId: String? = null,
+) =
     mock<LocalDataSource>().also { dataSource ->
         val driver = mock<DatabaseDriver>()
         `when`(driver.id).thenReturn("mongo")
@@ -450,8 +454,12 @@ internal fun mockDataSource(url: MongoDbServerUrl = MongoDbServerUrl("mongodb://
         `when`(dataSource.databaseDriver).thenReturn(driver)
         `when`(dataSource.dbms).thenReturn(Dbms.MONGO)
         val testClass = Thread.currentThread().stackTrace[2].className
-        `when`(dataSource.name).thenReturn(testClass + "_" + UUID.randomUUID().toString())
-        `when`(dataSource.uniqueId).thenReturn(testClass + "_" + UUID.randomUUID().toString())
+        `when`(dataSource.name).thenReturn(
+            name ?: (testClass + "_" + UUID.randomUUID().toString())
+        )
+        `when`(dataSource.uniqueId).thenReturn(
+            uniqueId ?: (testClass + "_" + UUID.randomUUID().toString())
+        )
     }
 
 /**

packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/ui/components/connection/ConnectionBootstrapCardTest.kt

@@ -14,6 +14,7 @@ import com.mongodb.jbplugin.fixtures.readClipboard
 import com.mongodb.jbplugin.fixtures.setContentWithTheme
 import com.mongodb.jbplugin.ui.viewModel.ConnectionState
 import com.mongodb.jbplugin.ui.viewModel.DatabaseState
+import com.mongodb.jbplugin.ui.viewModel.DatabasesLoadingState
 import com.mongodb.jbplugin.ui.viewModel.SelectedConnectionState
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertTrue
@@ -237,4 +238,106 @@ class ConnectionBootstrapCardTest {
         onNodeWithTag("DisconnectItem").performClick()
         assertTrue(clicked)
     }
+
+    @Test
+    fun `shows selected database when one is selected`() = runComposeUiTest {
+        val dataSource = mockDataSource()
+
+        setContentWithTheme {
+            _ConnectionBootstrapCard(
+                ConnectionState(
+                    listOf(dataSource),
+                    dataSource,
+                    SelectedConnectionState.Connected(dataSource)
+                ),
+                DatabaseState(
+                    databasesLoadingState = DatabasesLoadingState.Loaded(dataSource, listOf("testDB")),
+                    selectedDatabase = "testDB"
+                )
+            )
+        }
+
+        onNodeWithTag("DatabaseComboBox").assertTextEquals("testDB")
+    }
+
+    @Test
+    fun `shows loading databases state when fetching databases`() = runComposeUiTest {
+        val dataSource = mockDataSource()
+
+        setContentWithTheme {
+            _ConnectionBootstrapCard(
+                ConnectionState(
+                    listOf(dataSource),
+                    dataSource,
+                    SelectedConnectionState.Connected(dataSource)
+                ),
+                DatabaseState(
+                    databasesLoadingState = DatabasesLoadingState.Loading(dataSource),
+                    selectedDatabase = null
+                )
+            )
+        }
+
+        onNodeWithTag("DatabaseComboBox").assertTextEquals("Loading databases...")
+    }
+
+    @Test
+    fun `handles database selection through UI`() = runComposeUiTest {
+        val dataSource = mockDataSource()
+        var selectedDatabase: String? = null
+
+        setContentWithTheme {
+            CompositionLocalProvider(
+                LocalDatabaseCallbacks provides DatabaseCallbacks(
+                    selectDatabase = { selectedDatabase = it }
+                )
+            ) {
+                _ConnectionBootstrapCard(
+                    ConnectionState(
+                        listOf(dataSource),
+                        dataSource,
+                        SelectedConnectionState.Connected(dataSource)
+                    ),
+                    DatabaseState(
+                        databasesLoadingState = DatabasesLoadingState.Loaded(dataSource, listOf("testDB")),
+                        selectedDatabase = null
+                    )
+                )
+            }
+        }
+
+        onNodeWithTag("DatabaseComboBox").performClick()
+        onNodeWithTag("DatabaseItem::testDB").performClick()
+        assertEquals("testDB", selectedDatabase)
+    }
+
+    @Test
+    fun `handles database unselection through UI`() = runComposeUiTest {
+        val dataSource = mockDataSource()
+        var unselectClicked = false
+
+        setContentWithTheme {
+            CompositionLocalProvider(
+                LocalDatabaseCallbacks provides DatabaseCallbacks(
+                    unselectSelectedDatabase = { unselectClicked = true }
+                )
+            ) {
+                _ConnectionBootstrapCard(
+                    ConnectionState(
+                        listOf(dataSource),
+                        dataSource,
+                        SelectedConnectionState.Connected(dataSource)
+                    ),
+                    DatabaseState(
+                        databasesLoadingState = DatabasesLoadingState.Loaded(dataSource, listOf("testDB")),
+                        selectedDatabase = "testDB"
+                    )
+                )
+            }
+        }
+
+        onNodeWithTag("DatabaseComboBox").performClick()
+        onNodeWithTag("UnselectItem").performClick()
+        assertTrue(unselectClicked)
+    }
 }

packages/jetbrains-plugin/src/test/kotlin/com/mongodb/jbplugin/ui/viewModel/ConnectionStateViewModelTest.kt

@@ -16,6 +16,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertFalse
 import org.junit.jupiter.api.Assertions.assertTrue
@@ -421,4 +422,137 @@ class ConnectionStateViewModelTest {
             assertInstanceOf<DatabasesLoadingState.Failed>(viewModel.databaseState.value.databasesLoadingState)
         }
     }
+
+    @Test
+    fun `when a data source is changed, it updates the connection list and reconnects if selected`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        viewModel.connectionSaga = mock()
+        val dataSourceManager = mock<DataSourceManager<LocalDataSource>>()
+        val updatedDataSource = mockDataSource(name = dataSource.name, uniqueId = dataSource.uniqueId)
+
+        runBlocking {
+            viewModel.selectConnection(dataSource)
+            assertEquals(dataSource, viewModel.connectionState.value.selectedConnection)
+        }
+
+        viewModel.dataSourceChanged(dataSourceManager, updatedDataSource)
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            assertTrue(viewModel.connectionState.value.connections.any { it.uniqueId == updatedDataSource.uniqueId })
+            assertEquals(updatedDataSource, viewModel.connectionState.value.selectedConnection)
+            verify(viewModel.connectionSaga, timeout(1000)).connect(updatedDataSource)
+        }
+    }
+
+    @Test
+    fun `when selecting a database, it updates preferences and triggers editor reanalysis`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        viewModel.connectionSaga = mock()
+
+        runBlocking {
+            viewModel.selectDatabase("testDB")
+        }
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            assertEquals("testDB", viewModel.databaseState.value.selectedDatabase)
+            assertEquals("testDB", connectionPreferences.database)
+            verify(editorService, timeout(1000)).reAnalyzeSelectedEditor(true)
+        }
+    }
+
+    @Test
+    fun `when unselecting a database, it clears preferences and triggers editor reanalysis`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        viewModel.connectionSaga = mock()
+
+        runBlocking {
+            viewModel.selectDatabase("testDB")
+            viewModel.unselectSelectedDatabase()
+        }
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            assertEquals(null, viewModel.databaseState.value.selectedDatabase)
+            assertEquals(null, connectionPreferences.database)
+            verify(editorService, timeout(1000).times(2)).reAnalyzeSelectedEditor(true)
+        }
+    }
+
+    @Test
+    fun `when adding a new connection, it selects the connection if successfully created`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) = runTest {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        val newDataSource = mockDataSource()
+        whenever(connectionSaga.addNewConnection()).thenReturn(newDataSource)
+        viewModel.connectionSaga = connectionSaga
+
+        runBlocking {
+            viewModel.addNewConnection()
+        }
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            assertEquals(newDataSource, viewModel.connectionState.value.selectedConnection)
+            verify(connectionSaga, timeout(1000)).connect(newDataSource)
+        }
+    }
+
+    @Test
+    fun `when connection fails, it updates state with error message`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        val realConnectionSaga = viewModel.connectionSaga
+        val errorMessage = "Connection failed"
+        whenever(connectionSaga.connect(dataSource)).then {
+            coroutineScope.launch {
+                realConnectionSaga.emitSelectedConnectionStateChange(
+                    SelectedConnectionState.Failed(dataSource, errorMessage)
+                )
+            }
+        }
+        viewModel.connectionSaga = connectionSaga
+
+        runBlocking {
+            viewModel.selectConnection(dataSource)
+        }
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            val state = viewModel.connectionState.value.selectedConnectionState
+            assertInstanceOf<SelectedConnectionState.Failed>(state)
+            assertEquals(errorMessage, (state as SelectedConnectionState.Failed).errorMessage)
+        }
+    }
+
+    @Test
+    fun `when tool window is shown multiple times, it only loads initial state once`(
+        project: Project,
+        coroutineScope: TestScope,
+    ) {
+        val viewModel = ConnectionStateViewModel(project, coroutineScope)
+        viewModel.connectionSaga = connectionSaga
+
+        viewModel.toolWindowShown(mongoDBToolWindow)
+        viewModel.toolWindowShown(mongoDBToolWindow)
+
+        eventually {
+            coroutineScope.advanceUntilIdle()
+            verify(connectionSaga, timeout(1000).times(1)).listMongoDbConnections()
+        }
+    }
 }