Preserve selection when workspaces starts or stops

When a workspace starts the workspace row is replaced with a row for the
agent(s) and you have to select it again which is annoying (similarly
for when the workspace stops).  This preserves the selection when the
workspace transitions.

src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt

@@ -7,10 +7,16 @@ import com.coder.gateway.sdk.v2.models.WorkspaceTransition
 import java.util.UUID
 import javax.swing.Icon
 
+// TODO: Refactor to have a list of workspaces that each have agents.  We
+// present in the UI as a single flat list in the table (when there are no
+// agents we display a row for the workspace) but still, a list of workspaces
+// each with a list of agents might reflect reality more closely.  When we
+// iterate over the list we can add the workspace row if it has no agents
+// otherwise iterate over the agents and then flatten the result.
 data class WorkspaceAgentModel(
     val workspaceID: UUID,
     val workspaceName: String,
-    val name: String,
+    val name: String, // Name of the workspace OR the agent if this is for an agent.
     val templateID: UUID,
     val templateName: String,
     val templateIconPath: String,

src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt

@@ -669,11 +669,9 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
             }
         }
         withContext(Dispatchers.Main) {
-            val selectedWorkspace = tableOfWorkspaces.selectedObject?.name
+            val selectedWorkspace = tableOfWorkspaces.selectedObject
             tableOfWorkspaces.listTableModel.items = ws.toList()
-            if (selectedWorkspace != null) {
-                tableOfWorkspaces.selectItem(selectedWorkspace)
-            }
+            tableOfWorkspaces.selectItem(selectedWorkspace)
         }
     }
 
@@ -923,15 +921,31 @@ class WorkspacesTableModel : ListTableModel<WorkspaceAgentModel>(
 }
 
 class WorkspacesTable : TableView<WorkspaceAgentModel>(WorkspacesTableModel()) {
-    fun selectItem(workspaceName: String?) {
-        if (workspaceName != null) {
-            this.items.forEachIndexed { index, workspaceAgentModel ->
-                if (workspaceAgentModel.name == workspaceName) {
-                    selectionModel.addSelectionInterval(convertRowIndexToView(index), convertRowIndexToView(index))
-                    // fix cell selection case
-                    columnModel.selectionModel.addSelectionInterval(0, columnCount - 1)
-                }
-            }
+    /**
+     * Given either a workspace or an agent select in order of preference:
+     * 1. That same agent or workspace.
+     * 2. The first match for the workspace (workspace itself or first agent).
+     */
+    fun selectItem(workspace: WorkspaceAgentModel?) {
+        val index = getNewSelection(workspace)
+        if (index > -1) {
+            selectionModel.addSelectionInterval(convertRowIndexToView(index), convertRowIndexToView(index))
+            // Fix cell selection case.
+            columnModel.selectionModel.addSelectionInterval(0, columnCount - 1)
         }
     }
+
+    private fun getNewSelection(oldSelection: WorkspaceAgentModel?): Int {
+        if (oldSelection == null) {
+            return -1
+        }
+        val index = listTableModel.items.indexOfFirst {
+            it.name == oldSelection.name && it.workspaceName == oldSelection.workspaceName
+        }
+        if (index > -1) {
+            return index
+        }
+        return listTableModel.items.indexOfFirst { it.workspaceName == oldSelection.workspaceName }
+    }
+
 }

src/test/groovy/CoderWorkspacesStepViewTest.groovy

@@ -0,0 +1,54 @@
+import com.coder.gateway.views.steps.WorkspacesTable
+import spock.lang.Specification
+import spock.lang.Unroll
+
+@Unroll
+class CoderWorkspacesStepViewTest extends Specification {
+    def "gets new selection"() {
+        given:
+        def table = new WorkspacesTable()
+        table.listTableModel.items = List.of(
+                // An off workspace.
+                DataGen.workspace("ws1", "ws1"),
+
+                // On workspaces.
+                DataGen.workspace("agent1", "ws2"),
+                DataGen.workspace("agent2", "ws2"),
+                DataGen.workspace("agent3", "ws3"),
+
+                // Another off workspace.
+                DataGen.workspace("ws4", "ws4"),
+
+                // In practice we do not list both agents and workspaces
+                // together but here test that anyway with an agent first and
+                // then with a workspace first.
+                DataGen.workspace("agent2", "ws5"),
+                DataGen.workspace("ws5", "ws5"),
+                DataGen.workspace("ws6", "ws6"),
+                DataGen.workspace("agent3", "ws6"),
+        )
+
+        expect:
+        table.getNewSelection(selected) == expected
+
+        where:
+        selected                           | expected
+        null                               | -1 // No selection.
+        DataGen.workspace("gone", "gone")  | -1 // No workspace that matches.
+        DataGen.workspace("ws1", "ws1")    | 0  // Workspace exact match.
+        DataGen.workspace("gone", "ws1")   | 0  // Agent gone, select workspace.
+        DataGen.workspace("ws2", "ws2")    | 1  // Workspace gone, select first agent.
+        DataGen.workspace("agent1", "ws2") | 1  // Agent exact match.
+        DataGen.workspace("agent2", "ws2") | 2  // Agent exact match.
+        DataGen.workspace("ws3", "ws3")    | 3  // Workspace gone, select first agent.
+        DataGen.workspace("agent3", "ws3") | 3  // Agent exact match.
+        DataGen.workspace("gone", "ws4")   | 4  // Agent gone, select workspace.
+        DataGen.workspace("ws4", "ws4")    | 4  // Workspace exact match.
+        DataGen.workspace("agent2", "ws5") | 5  // Agent exact match.
+        DataGen.workspace("gone", "ws5")   | 5  // Agent gone, another agent comes first.
+        DataGen.workspace("ws5", "ws5")    | 6  // Workspace exact match.
+        DataGen.workspace("ws6", "ws6")    | 7  // Workspace exact match.
+        DataGen.workspace("gone", "ws6")   | 7  // Agent gone, workspace comes first.
+        DataGen.workspace("agent3", "ws6") | 8  // Agent exact match.
+    }
+}