From 4fd65ad1b5d457275bf32e364ec1f22f4bd2b16d Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Thu, 12 Oct 2023 13:24:44 -0800
Subject: [PATCH 01/10] Add method to get agents even for off workspaces

---
 .../gateway/sdk/CoderRestClientService.kt     | 19 +++++++++++++++++++
 .../coder/gateway/sdk/v2/CoderV2RestFacade.kt |  6 +++++-
 .../coder/gateway/sdk/v2/models/Workspace.kt  |  4 ++--
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
index 21ccf5fd..7c874ca3 100644
--- a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
+++ b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
@@ -1,5 +1,6 @@
 package com.coder.gateway.sdk
 
+import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.convertors.InstantConverter
 import com.coder.gateway.sdk.ex.AuthenticationResponseException
 import com.coder.gateway.sdk.ex.TemplateResponseException
@@ -12,6 +13,7 @@ import com.coder.gateway.sdk.v2.models.User
 import com.coder.gateway.sdk.v2.models.Workspace
 import com.coder.gateway.sdk.v2.models.WorkspaceBuild
 import com.coder.gateway.sdk.v2.models.WorkspaceTransition
+import com.coder.gateway.sdk.v2.models.toAgentModels
 import com.google.gson.Gson
 import com.google.gson.GsonBuilder
 import com.intellij.ide.plugins.PluginManagerCore
@@ -105,6 +107,23 @@ class CoderRestClient(var url: URL, var token: String, var headerCommand: String
         return workspacesResponse.body()!!.workspaces
     }
 
+    /**
+     * Retrieves agents for the specified workspaces.  Since the workspaces
+     * response does not include agents when the workspace is off, this fires
+     * off separate queries to get the agents for each workspace, just like
+     * `coder config-ssh` does (otherwise we risk removing hosts from the SSH
+     * config when they are off).
+     */
+    fun agents(workspaces: List<Workspace>): List<WorkspaceAgentModel> {
+        return workspaces.flatMap {
+            val resourcesResponse = retroRestClient.templateVersionResources(it.latestBuild.templateVersionID).execute()
+            if (!resourcesResponse.isSuccessful) {
+                throw WorkspaceResponseException("Unable to retrieve template resources for ${it.name} from $url: code ${resourcesResponse.code()}, reason: ${resourcesResponse.message().ifBlank { "no reason provided" }}")
+            }
+            it.toAgentModels(resourcesResponse.body()!!)
+        }
+    }
+
     fun buildInfo(): BuildInfo {
         val buildInfoResponse = retroRestClient.buildInfo().execute()
         if (!buildInfoResponse.isSuccessful) {
diff --git a/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt b/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt
index 036c25c3..06f1fc80 100644
--- a/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt
+++ b/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt
@@ -5,6 +5,7 @@ import com.coder.gateway.sdk.v2.models.CreateWorkspaceBuildRequest
 import com.coder.gateway.sdk.v2.models.Template
 import com.coder.gateway.sdk.v2.models.User
 import com.coder.gateway.sdk.v2.models.WorkspaceBuild
+import com.coder.gateway.sdk.v2.models.WorkspaceResource
 import com.coder.gateway.sdk.v2.models.WorkspacesResponse
 import retrofit2.Call
 import retrofit2.http.Body
@@ -39,4 +40,7 @@ interface CoderV2RestFacade {
 
     @GET("api/v2/templates/{templateID}")
     fun template(@Path("templateID") templateID: UUID): Call<Template>
-}
\ No newline at end of file
+
+    @GET("api/v2/templateversions/{templateID}/resources")
+    fun templateVersionResources(@Path("templateID") templateID: UUID): Call<List<WorkspaceResource>>
+}
diff --git a/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt b/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
index 129489c6..ca0bd5f0 100644
--- a/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
+++ b/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
@@ -31,8 +31,8 @@ data class Workspace(
     @SerializedName("last_used_at") val lastUsedAt: Instant,
 )
 
-fun Workspace.toAgentModels(): Set<WorkspaceAgentModel> {
-    val wam = this.latestBuild.resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
+fun Workspace.toAgentModels(resources: List<WorkspaceResource> = this.latestBuild.resources): Set<WorkspaceAgentModel> {
+    val wam = resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
         val workspaceWithAgentName = "${this.name}.${agent.name}"
         val wm = WorkspaceAgentModel(
             this.id,

From d41c1fd923abda3f7f8b567d12cc6ad88480e555 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Thu, 12 Oct 2023 13:28:19 -0800
Subject: [PATCH 02/10] Get all agents when configuring SSH

---
 .../kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt | 2 +-
 .../com/coder/gateway/views/steps/CoderWorkspacesStepView.kt   | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
index f685ebc5..91cf7ece 100644
--- a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
+++ b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -100,7 +100,7 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
             cli.login(client.token)
 
             indicator.text = "Configuring Coder CLI..."
-            cli.configSsh(workspaces.flatMap { it.toAgentModels() }, settings.headerCommand)
+            cli.configSsh(client.agents(workspaces), settings.headerCommand)
 
             // TODO: Ask for these if missing.  Maybe we can reuse the second
             //  step of the wizard?  Could also be nice if we automatically used
diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
index 1d7da3ff..2ddefd4e 100644
--- a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
+++ b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
@@ -620,7 +620,8 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
         poller?.cancel()
 
         logger.info("Configuring Coder CLI...")
-        cli.configSsh(tableOfWorkspaces.items, settings.headerCommand)
+        val workspaces = clientService.client.workspaces()
+        cli.configSsh(clientService.client.agents(workspaces), settings.headerCommand)
 
         // The config directory can be used to pull the URL and token in
         // order to query this workspace's status in other flows, for

From feb55acf4f6b1e9a97ba0430903d08432dfee117 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Fri, 13 Oct 2023 14:43:44 -0800
Subject: [PATCH 03/10] Support agent ID parameter

---
 .../gateway/CoderGatewayConnectionProvider.kt | 20 ++++++++++++++-----
 .../gateway/models/WorkspaceAgentModel.kt     |  1 +
 .../coder/gateway/sdk/v2/models/Workspace.kt  |  2 ++
 src/test/groovy/DataGen.groovy                |  1 +
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
index 91cf7ece..1467a13e 100644
--- a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
+++ b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -23,7 +23,8 @@ import java.net.URL
 private const val URL = "url"
 private const val TOKEN = "token"
 private const val WORKSPACE = "workspace"
-private const val AGENT = "agent"
+private const val AGENT_NAME = "agent"
+private const val AGENT_ID = "agent_id"
 private const val FOLDER = "folder"
 private const val IDE_DOWNLOAD_LINK = "ide_download_link"
 private const val IDE_PRODUCT_CODE = "ide_product_code"
@@ -72,14 +73,23 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
             }
 
             // If the agent is missing and the workspace has only one, use that.
-            val agent = if (!parameters[AGENT].isNullOrBlank())
-                agents.firstOrNull { it.name == "$workspaceName.${parameters[AGENT]}"}
+            // Prefer the ID over the name if both are set.
+            val agent = if (!parameters[AGENT_ID].isNullOrBlank())
+                agents.firstOrNull {it.agentID.toString() == parameters[AGENT_ID]}
+            else if (!parameters[AGENT_NAME].isNullOrBlank())
+                agents.firstOrNull { it.name == "$workspaceName.${parameters[AGENT_NAME]}"}
             else if (agents.size == 1) agents.first()
             else null
 
             if (agent == null) {
-                // TODO: Show a dropdown and ask for an agent.
-                throw IllegalArgumentException("Query parameter \"$AGENT\" is missing")
+                if (parameters[AGENT_ID].isNullOrBlank() && parameters[AGENT_NAME].isNullOrBlank()) {
+                    // TODO: Show a dropdown and ask for an agent.
+                    throw IllegalArgumentException("Unable to determine which agent to connect to; one of \"$AGENT_NAME\" or \"$AGENT_ID\" must be set because \"$workspaceName\" has more than one agent")
+                } else if (parameters[AGENT_ID].isNullOrBlank()) {
+                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
+                } else {
+                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent named \"${parameters[AGENT_NAME]}\"")
+                }
             }
 
             if (agent.agentStatus.pending()) {
diff --git a/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt b/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
index c4893b9c..61276100 100644
--- a/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
+++ b/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
@@ -14,6 +14,7 @@ import javax.swing.Icon
 // 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 agentID: UUID?,
     val workspaceID: UUID,
     val workspaceName: String,
     val name: String, // Name of the workspace OR the agent if this is for an agent.
diff --git a/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt b/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
index ca0bd5f0..89ce53c7 100644
--- a/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
+++ b/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
@@ -35,6 +35,7 @@ fun Workspace.toAgentModels(resources: List<WorkspaceResource> = this.latestBuil
     val wam = resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
         val workspaceWithAgentName = "${this.name}.${agent.name}"
         val wm = WorkspaceAgentModel(
+            agent.id,
             this.id,
             this.name,
             workspaceWithAgentName,
@@ -55,6 +56,7 @@ fun Workspace.toAgentModels(resources: List<WorkspaceResource> = this.latestBuil
     }.toSet()
     if (wam.isNullOrEmpty()) {
         val wm = WorkspaceAgentModel(
+            null,
             this.id,
             this.name,
             this.name,
diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
index af609f99..69d329fe 100644
--- a/src/test/groovy/DataGen.groovy
+++ b/src/test/groovy/DataGen.groovy
@@ -7,6 +7,7 @@ import com.coder.gateway.sdk.v2.models.WorkspaceTransition
 class DataGen {
     static WorkspaceAgentModel workspace(String name, String workspaceName = name) {
         return new WorkspaceAgentModel(
+                UUID.randomUUID(),
                 UUID.randomUUID(),
                 workspaceName,
                 name,

From a0f22398d7776e6f83f1f4eb71f29cafce84d76a Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Mon, 16 Oct 2023 11:19:51 -0800
Subject: [PATCH 04/10] Fix DataGen not quite mirroring toAgentModels()

It should omit agent information when the workspace specifies no agents
and the name should combine the workspace and agent names.
---
 .../coder/gateway/models/WorkspaceAgentModel.kt   |  2 +-
 src/test/groovy/DataGen.groovy                    | 15 +++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt b/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
index 61276100..d9678422 100644
--- a/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
+++ b/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
@@ -17,7 +17,7 @@ data class WorkspaceAgentModel(
     val agentID: UUID?,
     val workspaceID: UUID,
     val workspaceName: String,
-    val name: String, // Name of the workspace OR the agent if this is for an agent.
+    val name: String, // Name of the workspace OR workspace.agent if this is for an agent.
     val templateID: UUID,
     val templateName: String,
     val templateIconPath: String,
diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
index 69d329fe..206909f2 100644
--- a/src/test/groovy/DataGen.groovy
+++ b/src/test/groovy/DataGen.groovy
@@ -5,12 +5,19 @@ import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.sdk.v2.models.WorkspaceTransition
 
 class DataGen {
-    static WorkspaceAgentModel workspace(String name, String workspaceName = name) {
+    // Create a random workspace agent model.  If the workspace name is omitted
+    // then return a model without any agent bits, similar to what
+    // toAgentModels() does if the workspace does not specify any agents.
+    // TODO: Maybe better to randomly generate the workspace and then call
+    //       toAgentModels() on it.  Also the way an "agent" model can have no
+    //       agent in it seems weird; can we refactor to remove
+    //       WorkspaceAgentModel and use the original structs from the API?
+    static WorkspaceAgentModel workspace(String name, String workspaceName = "", UUID agentId = UUID.randomUUID()) {
         return new WorkspaceAgentModel(
+                workspaceName == "" ? null : agentId,
                 UUID.randomUUID(),
-                UUID.randomUUID(),
-                workspaceName,
-                name,
+                workspaceName == "" ? name : workspaceName,
+                workspaceName == "" ? name : (workspaceName + "." + name),
                 UUID.randomUUID(),
                 "template-name",
                 "template-icon-path",

From 800313029213afe5229267443cec5d65cbf20212 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Mon, 16 Oct 2023 11:22:18 -0800
Subject: [PATCH 05/10] Update workspace generator function name

This generates agent models, not workspaces.  Renaming opens the way for
generating random workspaces.
---
 src/test/groovy/CoderCLIManagerTest.groovy    |  4 +-
 .../groovy/CoderWorkspacesStepViewTest.groovy | 54 +++++++++----------
 src/test/groovy/DataGen.groovy                |  2 +-
 3 files changed, 30 insertions(+), 30 deletions(-)

diff --git a/src/test/groovy/CoderCLIManagerTest.groovy b/src/test/groovy/CoderCLIManagerTest.groovy
index effa5605..c6b7e1a6 100644
--- a/src/test/groovy/CoderCLIManagerTest.groovy
+++ b/src/test/groovy/CoderCLIManagerTest.groovy
@@ -413,7 +413,7 @@ class CoderCLIManagerTest extends Specification {
                 .replace("/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64", CoderCLIManager.escape(ccm.localBinaryPath.toString()))
 
         when:
-        ccm.configSsh(workspaces.collect { DataGen.workspace(it) }, headerCommand)
+        ccm.configSsh(workspaces.collect { DataGen.workspaceAgentModel(it) }, headerCommand)
 
         then:
         sshConfigPath.toFile().text == expectedConf
@@ -473,7 +473,7 @@ class CoderCLIManagerTest extends Specification {
         def ccm = new CoderCLIManager(new URL("https://test.coder.invalid"), tmpdir)
 
         when:
-        ccm.configSsh(["foo", "bar"].collect { DataGen.workspace(it) }, headerCommand)
+        ccm.configSsh(["foo", "bar"].collect { DataGen.workspaceAgentModel(it) }, headerCommand)
 
         then:
         thrown(Exception)
diff --git a/src/test/groovy/CoderWorkspacesStepViewTest.groovy b/src/test/groovy/CoderWorkspacesStepViewTest.groovy
index c684e963..d07dde3d 100644
--- a/src/test/groovy/CoderWorkspacesStepViewTest.groovy
+++ b/src/test/groovy/CoderWorkspacesStepViewTest.groovy
@@ -9,46 +9,46 @@ class CoderWorkspacesStepViewTest extends Specification {
         def table = new WorkspacesTable()
         table.listTableModel.items = List.of(
                 // An off workspace.
-                DataGen.workspace("ws1", "ws1"),
+                DataGen.workspaceAgentModel("ws1", "ws1"),
 
                 // On workspaces.
-                DataGen.workspace("agent1", "ws2"),
-                DataGen.workspace("agent2", "ws2"),
-                DataGen.workspace("agent3", "ws3"),
+                DataGen.workspaceAgentModel("agent1", "ws2"),
+                DataGen.workspaceAgentModel("agent2", "ws2"),
+                DataGen.workspaceAgentModel("agent3", "ws3"),
 
                 // Another off workspace.
-                DataGen.workspace("ws4", "ws4"),
+                DataGen.workspaceAgentModel("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"),
+                DataGen.workspaceAgentModel("agent2", "ws5"),
+                DataGen.workspaceAgentModel("ws5", "ws5"),
+                DataGen.workspaceAgentModel("ws6", "ws6"),
+                DataGen.workspaceAgentModel("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.
+        selected                                     | expected
+        null                                         | -1 // No selection.
+        DataGen.workspaceAgentModel("gone", "gone")  | -1 // No workspace that matches.
+        DataGen.workspaceAgentModel("ws1", "ws1")    | 0  // Workspace exact match.
+        DataGen.workspaceAgentModel("gone", "ws1")   | 0  // Agent gone, select workspace.
+        DataGen.workspaceAgentModel("ws2", "ws2")    | 1  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("agent1", "ws2") | 1  // Agent exact match.
+        DataGen.workspaceAgentModel("agent2", "ws2") | 2  // Agent exact match.
+        DataGen.workspaceAgentModel("ws3", "ws3")    | 3  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("agent3", "ws3") | 3  // Agent exact match.
+        DataGen.workspaceAgentModel("gone", "ws4")   | 4  // Agent gone, select workspace.
+        DataGen.workspaceAgentModel("ws4", "ws4")    | 4  // Workspace exact match.
+        DataGen.workspaceAgentModel("agent2", "ws5") | 5  // Agent exact match.
+        DataGen.workspaceAgentModel("gone", "ws5")   | 5  // Agent gone, another agent comes first.
+        DataGen.workspaceAgentModel("ws5", "ws5")    | 6  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws6", "ws6")    | 7  // Workspace exact match.
+        DataGen.workspaceAgentModel("gone", "ws6")   | 7  // Agent gone, workspace comes first.
+        DataGen.workspaceAgentModel("agent3", "ws6") | 8  // Agent exact match.
     }
 }
diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
index 206909f2..71e8e060 100644
--- a/src/test/groovy/DataGen.groovy
+++ b/src/test/groovy/DataGen.groovy
@@ -12,7 +12,7 @@ class DataGen {
     //       toAgentModels() on it.  Also the way an "agent" model can have no
     //       agent in it seems weird; can we refactor to remove
     //       WorkspaceAgentModel and use the original structs from the API?
-    static WorkspaceAgentModel workspace(String name, String workspaceName = "", UUID agentId = UUID.randomUUID()) {
+    static WorkspaceAgentModel workspaceAgentModel(String name, String workspaceName = "", UUID agentId = UUID.randomUUID()) {
         return new WorkspaceAgentModel(
                 workspaceName == "" ? null : agentId,
                 UUID.randomUUID(),

From e3d4f7f4c7b47ff73771b8f63f5d2f26fdf19a57 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Mon, 16 Oct 2023 11:37:42 -0800
Subject: [PATCH 06/10] Break out getting matching agents

I also made a tweak to check that the agent ID is not null since
toAgentModels() will return the workspace without any agent bits set if
there are no agents.

And the wrong error message would show when either the id or name were
missing.  I also flipped them around while fixing this to match the
order above it.
---
 .../gateway/CoderGatewayConnectionProvider.kt |  72 +++++++----
 .../CoderGatewayConnectionProviderTest.groovy | 114 ++++++++++++++++++
 src/test/groovy/DataGen.groovy                |  93 +++++++++++++-
 3 files changed, 253 insertions(+), 26 deletions(-)
 create mode 100644 src/test/groovy/CoderGatewayConnectionProviderTest.groovy

diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
index 1467a13e..99fdd9da 100644
--- a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
+++ b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -3,10 +3,12 @@
 package com.coder.gateway
 
 import com.coder.gateway.models.TokenSource
+import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.CoderCLIManager
 import com.coder.gateway.sdk.CoderRestClient
 import com.coder.gateway.sdk.ex.AuthenticationResponseException
 import com.coder.gateway.sdk.toURL
+import com.coder.gateway.sdk.v2.models.Workspace
 import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.sdk.v2.models.toAgentModels
 import com.coder.gateway.sdk.withPath
@@ -67,30 +69,8 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
                 WorkspaceStatus.RUNNING -> Unit // All is well
             }
 
-            val agents = workspace.toAgentModels()
-            if (agents.isEmpty()) {
-                throw IllegalArgumentException("The workspace \"$workspaceName\" has no agents")
-            }
-
-            // If the agent is missing and the workspace has only one, use that.
-            // Prefer the ID over the name if both are set.
-            val agent = if (!parameters[AGENT_ID].isNullOrBlank())
-                agents.firstOrNull {it.agentID.toString() == parameters[AGENT_ID]}
-            else if (!parameters[AGENT_NAME].isNullOrBlank())
-                agents.firstOrNull { it.name == "$workspaceName.${parameters[AGENT_NAME]}"}
-            else if (agents.size == 1) agents.first()
-            else null
-
-            if (agent == null) {
-                if (parameters[AGENT_ID].isNullOrBlank() && parameters[AGENT_NAME].isNullOrBlank()) {
-                    // TODO: Show a dropdown and ask for an agent.
-                    throw IllegalArgumentException("Unable to determine which agent to connect to; one of \"$AGENT_NAME\" or \"$AGENT_ID\" must be set because \"$workspaceName\" has more than one agent")
-                } else if (parameters[AGENT_ID].isNullOrBlank()) {
-                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
-                } else {
-                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent named \"${parameters[AGENT_NAME]}\"")
-                }
-            }
+            // TODO: Show a dropdown and ask for an agent if missing.
+            val agent = getMatchingAgent(parameters, workspace)
 
             if (agent.agentStatus.pending()) {
                 // TODO: Wait for the agent to be ready.
@@ -211,5 +191,49 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
 
     companion object {
         val logger = Logger.getInstance(CoderGatewayConnectionProvider::class.java.simpleName)
+
+        /**
+         * Return the agent matching the provided agent ID or name in the
+         * parameters.  The name is ignored if the ID is set.  If neither was
+         * supplied and the workspace has only one agent, return that.
+         * Otherwise throw an error.
+         *
+         * @throws [MissingArgumentException, IllegalArgumentException]
+         */
+        @JvmStatic
+        fun getMatchingAgent(parameters: Map<String, String>, workspace: Workspace): WorkspaceAgentModel {
+            // A WorkspaceAgentModel will still be returned if there are no
+            // agents; in this case it represents the workspace instead.
+            // TODO: Seems confusing for something with "agent" in the name to
+            //       potentially not actually be an agent; can we replace
+            //       WorkspaceAgentModel with the original structs from the API?
+            val agents = workspace.toAgentModels()
+            if (agents.isEmpty() || (agents.size == 1 && agents.first().agentID == null)) {
+                throw IllegalArgumentException("The workspace \"${workspace.name}\" has no agents")
+            }
+
+            // If the agent is missing and the workspace has only one, use that.
+            // Prefer the ID over the name if both are set.
+            val agent = if (!parameters[AGENT_ID].isNullOrBlank())
+                agents.firstOrNull { it.agentID.toString() == parameters[AGENT_ID] }
+            else if (!parameters[AGENT_NAME].isNullOrBlank())
+                agents.firstOrNull { it.name == "${workspace.name}.${parameters[AGENT_NAME]}"}
+            else if (agents.size == 1) agents.first()
+            else null
+
+            if (agent == null) {
+                if (!parameters[AGENT_ID].isNullOrBlank()) {
+                    throw IllegalArgumentException("The workspace \"${workspace.name}\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
+                } else if (!parameters[AGENT_NAME].isNullOrBlank()){
+                    throw IllegalArgumentException("The workspace \"${workspace.name}\"does not have an agent named \"${parameters[AGENT_NAME]}\"")
+                } else {
+                    throw MissingArgumentException("Unable to determine which agent to connect to; one of \"$AGENT_NAME\" or \"$AGENT_ID\" must be set because the workspace \"${workspace.name}\" has more than one agent")
+                }
+            }
+
+            return agent
+        }
     }
 }
+
+class MissingArgumentException(message: String) : IllegalArgumentException(message)
diff --git a/src/test/groovy/CoderGatewayConnectionProviderTest.groovy b/src/test/groovy/CoderGatewayConnectionProviderTest.groovy
new file mode 100644
index 00000000..5d5008ff
--- /dev/null
+++ b/src/test/groovy/CoderGatewayConnectionProviderTest.groovy
@@ -0,0 +1,114 @@
+package com.coder.gateway
+
+import spock.lang.Shared
+import spock.lang.Specification
+import spock.lang.Unroll
+
+@Unroll
+class CoderGatewayConnectionProviderTest extends Specification {
+    @Shared
+    def agents = [
+        agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24",
+        agent_name_2: "fb3daea4-da6b-424d-84c7-36b90574cfef",
+        agent_name: "9a920eee-47fb-4571-9501-e4b3120c12f2",
+    ]
+    def oneAgent = [
+        agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+    ]
+
+    def "gets matching agent"() {
+        expect:
+        def ws = DataGen.workspace("ws", agents)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString(expected)
+
+        where:
+        parameters                                         | expected
+        [agent:    "agent_name"]                           | "9a920eee-47fb-4571-9501-e4b3120c12f2"
+        [agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | "9a920eee-47fb-4571-9501-e4b3120c12f2"
+        [agent:    "agent_name_2"]                         | "fb3daea4-da6b-424d-84c7-36b90574cfef"
+        [agent_id: "fb3daea4-da6b-424d-84c7-36b90574cfef"] | "fb3daea4-da6b-424d-84c7-36b90574cfef"
+        [agent:    "agent_name_3"]                         | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+        [agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+
+        // Prefer agent_id.
+        [agent: "agent_name", agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+    }
+
+    def "fails to get matching agent"() {
+        when:
+        def ws = DataGen.workspace("ws", agents)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [:]                                                | MissingArgumentException | "Unable to determine"
+        [agent: ""]                                        | MissingArgumentException | "Unable to determine"
+        [agent_id: ""]                                     | MissingArgumentException | "Unable to determine"
+        [agent: null]                                      | MissingArgumentException | "Unable to determine"
+        [agent_id: null]                                   | MissingArgumentException | "Unable to determine"
+        [agent: "ws"]                                      | IllegalArgumentException | "agent named"
+        [agent: "ws.agent_name"]                           | IllegalArgumentException | "agent named"
+        [agent: "agent_name_4"]                            | IllegalArgumentException | "agent named"
+        [agent_id: "not-a-uuid"]                           | IllegalArgumentException | "agent with ID"
+        [agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+
+        // Will ignore agent if agent_id is set even if agent matches.
+        [agent: "agent_name", agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+    }
+
+    def "gets the first agent when workspace has only one"() {
+        expect:
+        def ws = DataGen.workspace("ws", oneAgent)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString("b0e4c54d-9ba9-4413-8512-11ca1e826a24")
+
+        where:
+        parameters << [
+            [:],
+            [agent: ""],
+            [agent_id: ""],
+            [agent: null],
+            [agent_id: null],
+        ]
+    }
+
+    def "fails to get agent when workspace has only one"() {
+        when:
+        def ws = DataGen.workspace("ws", oneAgent)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [agent:    "ws"]                                   | IllegalArgumentException | "agent named"
+        [agent:    "ws.agent_name_3"]                      | IllegalArgumentException | "agent named"
+        [agent:    "agent_name_4"]                         | IllegalArgumentException | "agent named"
+        [agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+    }
+
+    def "fails to get agent from workspace without agents"() {
+        when:
+        def ws = DataGen.workspace("ws")
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [:]                                                | IllegalArgumentException | "has no agents"
+        [agent: ""]                                        | IllegalArgumentException | "has no agents"
+        [agent_id: ""]                                     | IllegalArgumentException | "has no agents"
+        [agent: null]                                      | IllegalArgumentException | "has no agents"
+        [agent_id: null]                                   | IllegalArgumentException | "has no agents"
+        [agent:    "agent_name"]                           | IllegalArgumentException | "has no agents"
+        [agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | IllegalArgumentException | "has no agents"
+    }
+}
diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
index 71e8e060..62d8abd2 100644
--- a/src/test/groovy/DataGen.groovy
+++ b/src/test/groovy/DataGen.groovy
@@ -1,8 +1,7 @@
 import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.models.WorkspaceAndAgentStatus
 import com.coder.gateway.models.WorkspaceVersionStatus
-import com.coder.gateway.sdk.v2.models.WorkspaceStatus
-import com.coder.gateway.sdk.v2.models.WorkspaceTransition
+import com.coder.gateway.sdk.v2.models.*
 
 class DataGen {
     // Create a random workspace agent model.  If the workspace name is omitted
@@ -31,4 +30,94 @@ class DataGen {
                 null
         )
     }
+
+    static Workspace workspace(String name, Map<String, String> agents = [:]) {
+        UUID wsId = UUID.randomUUID()
+        UUID ownerId = UUID.randomUUID()
+        List<WorkspaceResource> resources = agents.collect{ agentName, agentId -> new WorkspaceResource(
+                UUID.randomUUID(),      // id
+                new Date().toInstant(), // created_at
+                UUID.randomUUID(),      // job_id
+                WorkspaceTransition.START,
+                "type",
+                "name",
+                false,                  // hide
+                "icon",
+                List.of(new WorkspaceAgent(
+                        UUID.fromString(agentId),
+                        new Date().toInstant(),    // created_at
+                        new Date().toInstant(),    // updated_at
+                        null,                      // first_connected_at
+                        null,                      // last_connected_at
+                        null,                      // disconnected_at
+                        WorkspaceAgentStatus.CONNECTED,
+                        agentName,
+                        UUID.randomUUID(),         // resource_id
+                        null,                      // instance_id
+                        "arch",                    // architecture
+                        [:],                       // environment_variables
+                        "os",                      // operating_system
+                        null,                      // startup_script
+                        null,                      // directory
+                        null,                      // expanded_directory
+                        "version",                 // version
+                        List.of(),                 // apps
+                        null,                      // latency
+                        0,                         // connection_timeout_seconds
+                        "url",                     // troubleshooting_url
+                        WorkspaceAgentLifecycleState.READY,
+                        false,                     // login_before_ready
+                )),
+                null,                   // metadata
+                0,                      // daily_cost
+        )}
+        return new Workspace(
+                wsId,
+                new Date().toInstant(), // created_at
+                new Date().toInstant(), // updated_at
+                ownerId,
+                "owner-name",
+                UUID.randomUUID(),      // template_id
+                "template-name",
+                "template-display-name",
+                "template-icon",
+                false,                  // template_allow_user_cancel_workspace_jobs
+                new WorkspaceBuild(
+                        UUID.randomUUID(),      // id
+                        new Date().toInstant(), // created_at
+                        new Date().toInstant(), // updated_at
+                        wsId,
+                        name,
+                        ownerId,
+                        "owner-name",
+                        UUID.randomUUID(),      // template_version_id
+                        0,                      // build_number
+                        WorkspaceTransition.START,
+                        UUID.randomUUID(),      // initiator_id
+                        "initiator-name",
+                        new ProvisionerJob(
+                                UUID.randomUUID(),      // id
+                                new Date().toInstant(), // created_at
+                                null,                   // started_at
+                                null,                   // completed_at
+                                null,                   // canceled_at
+                                null,                   // error
+                                ProvisionerJobStatus.SUCCEEDED,
+                                null,                   // worker_id
+                                UUID.randomUUID(),      // file_id
+                                [:],                    // tags
+                        ),
+                        BuildReason.INITIATOR,
+                        resources,
+                        null,                   // deadline
+                        WorkspaceStatus.RUNNING,
+                        0,                      // daily_cost
+                ),
+                false,                  // outdated
+                name,
+                null,                   // autostart_schedule
+                null,                   // ttl_ms
+                new Date().toInstant(), // last_used_at
+        )
+    }
 }

From 149536c180935fd58bc61aa1b69a843d64fb361f Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 17 Oct 2023 12:56:34 -0800
Subject: [PATCH 07/10] Make pluginVersion an constructor option

PluginManagerCore is throwing a null pointer exception when used in the
tests.
---
 .../coder/gateway/CoderGatewayConnectionProvider.kt |  2 +-
 .../com/coder/gateway/sdk/CoderRestClientService.kt | 13 +++++++++----
 .../CoderGatewayRecentWorkspaceConnectionsView.kt   |  2 +-
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
index 99fdd9da..e6acb61c 100644
--- a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
+++ b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -140,7 +140,7 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
         if (token == null) { // User aborted.
             throw IllegalArgumentException("Unable to connect to $deploymentURL, $TOKEN is missing")
         }
-        val client = CoderRestClient(deploymentURL, token.first, settings.headerCommand)
+        val client = CoderRestClient(deploymentURL, token.first, settings.headerCommand, null)
         return try {
             Pair(client, client.me().username)
         } catch (ex: AuthenticationResponseException) {
diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
index 7c874ca3..829d76fe 100644
--- a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
+++ b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
@@ -45,7 +45,7 @@ class CoderRestClientService {
      * @throws [AuthenticationResponseException] if authentication failed.
      */
     fun initClientSession(url: URL, token: String, headerCommand: String?): User {
-        client = CoderRestClient(url, token, headerCommand)
+        client = CoderRestClient(url, token, headerCommand, null)
         me = client.me()
         buildVersion = client.buildInfo().version
         isReady = true
@@ -53,17 +53,22 @@ class CoderRestClientService {
     }
 }
 
-class CoderRestClient(var url: URL, var token: String, var headerCommand: String?) {
+class CoderRestClient(var url: URL, var token: String,
+    private var headerCommand: String?,
+    private var pluginVersion: String?,
+) {
     private var httpClient: OkHttpClient
     private var retroRestClient: CoderV2RestFacade
 
     init {
         val gson: Gson = GsonBuilder().registerTypeAdapter(Instant::class.java, InstantConverter()).setPrettyPrinting().create()
-        val pluginVersion = PluginManagerCore.getPlugin(PluginId.getId("com.coder.gateway"))!! // this is the id from the plugin.xml
+        if (pluginVersion.isNullOrBlank()) {
+            pluginVersion = PluginManagerCore.getPlugin(PluginId.getId("com.coder.gateway"))!!.version // this is the id from the plugin.xml
+        }
 
         httpClient = OkHttpClient.Builder()
             .addInterceptor { it.proceed(it.request().newBuilder().addHeader("Coder-Session-Token", token).build()) }
-            .addInterceptor { it.proceed(it.request().newBuilder().addHeader("User-Agent", "Coder Gateway/${pluginVersion.version} (${SystemInfo.getOsNameAndVersion()}; ${SystemInfo.OS_ARCH})").build()) }
+            .addInterceptor { it.proceed(it.request().newBuilder().addHeader("User-Agent", "Coder Gateway/${pluginVersion} (${SystemInfo.getOsNameAndVersion()}; ${SystemInfo.OS_ARCH})").build()) }
             .addInterceptor {
                 var request = it.request()
                 val headers = getHeaders(url, headerCommand)
diff --git a/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt b/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt
index 1c7d0d71..4e5a5fcf 100644
--- a/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt
+++ b/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt
@@ -256,7 +256,7 @@ class CoderGatewayRecentWorkspaceConnectionsView(private val setContentCallback:
                 deployments[dir] ?: try {
                     val url = Path.of(dir).resolve("url").readText()
                     val token = Path.of(dir).resolve("session").readText()
-                    DeploymentInfo(CoderRestClient(url.toURL(), token, settings.headerCommand))
+                    DeploymentInfo(CoderRestClient(url.toURL(), token, settings.headerCommand, null))
                 } catch (e: Exception) {
                     logger.error("Unable to create client from $dir", e)
                     DeploymentInfo(error = "Error trying to read $dir: ${e.message}")

From a415c04b31990f0bd4e92c3f630ef41dfbc44d43 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 17 Oct 2023 13:51:46 -0800
Subject: [PATCH 08/10] Update off workspace generation

The second parameter should be blank to generate one without an
agent (which simulates an off workspace here).
---
 .../groovy/CoderWorkspacesStepViewTest.groovy | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/test/groovy/CoderWorkspacesStepViewTest.groovy b/src/test/groovy/CoderWorkspacesStepViewTest.groovy
index d07dde3d..883863f9 100644
--- a/src/test/groovy/CoderWorkspacesStepViewTest.groovy
+++ b/src/test/groovy/CoderWorkspacesStepViewTest.groovy
@@ -9,7 +9,7 @@ class CoderWorkspacesStepViewTest extends Specification {
         def table = new WorkspacesTable()
         table.listTableModel.items = List.of(
                 // An off workspace.
-                DataGen.workspaceAgentModel("ws1", "ws1"),
+                DataGen.workspaceAgentModel("ws1"),
 
                 // On workspaces.
                 DataGen.workspaceAgentModel("agent1", "ws2"),
@@ -17,14 +17,14 @@ class CoderWorkspacesStepViewTest extends Specification {
                 DataGen.workspaceAgentModel("agent3", "ws3"),
 
                 // Another off workspace.
-                DataGen.workspaceAgentModel("ws4", "ws4"),
+                DataGen.workspaceAgentModel("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.workspaceAgentModel("agent2", "ws5"),
-                DataGen.workspaceAgentModel("ws5", "ws5"),
-                DataGen.workspaceAgentModel("ws6", "ws6"),
+                DataGen.workspaceAgentModel("ws5"),
+                DataGen.workspaceAgentModel("ws6"),
                 DataGen.workspaceAgentModel("agent3", "ws6"),
         )
 
@@ -35,19 +35,19 @@ class CoderWorkspacesStepViewTest extends Specification {
         selected                                     | expected
         null                                         | -1 // No selection.
         DataGen.workspaceAgentModel("gone", "gone")  | -1 // No workspace that matches.
-        DataGen.workspaceAgentModel("ws1", "ws1")    | 0  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws1")           | 0  // Workspace exact match.
         DataGen.workspaceAgentModel("gone", "ws1")   | 0  // Agent gone, select workspace.
-        DataGen.workspaceAgentModel("ws2", "ws2")    | 1  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("ws2")           | 1  // Workspace gone, select first agent.
         DataGen.workspaceAgentModel("agent1", "ws2") | 1  // Agent exact match.
         DataGen.workspaceAgentModel("agent2", "ws2") | 2  // Agent exact match.
-        DataGen.workspaceAgentModel("ws3", "ws3")    | 3  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("ws3")           | 3  // Workspace gone, select first agent.
         DataGen.workspaceAgentModel("agent3", "ws3") | 3  // Agent exact match.
         DataGen.workspaceAgentModel("gone", "ws4")   | 4  // Agent gone, select workspace.
-        DataGen.workspaceAgentModel("ws4", "ws4")    | 4  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws4")           | 4  // Workspace exact match.
         DataGen.workspaceAgentModel("agent2", "ws5") | 5  // Agent exact match.
         DataGen.workspaceAgentModel("gone", "ws5")   | 5  // Agent gone, another agent comes first.
-        DataGen.workspaceAgentModel("ws5", "ws5")    | 6  // Workspace exact match.
-        DataGen.workspaceAgentModel("ws6", "ws6")    | 7  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws5")           | 6  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws6")           | 7  // Workspace exact match.
         DataGen.workspaceAgentModel("gone", "ws6")   | 7  // Agent gone, workspace comes first.
         DataGen.workspaceAgentModel("agent3", "ws6") | 8  // Agent exact match.
     }

From 1c6176733c307924d920741f6287b90f3dc6369e Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 17 Oct 2023 15:20:28 -0800
Subject: [PATCH 09/10] Break out random workspace resource

This will be used in a test.
---
 src/test/groovy/DataGen.groovy | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
index 62d8abd2..0025a8b0 100644
--- a/src/test/groovy/DataGen.groovy
+++ b/src/test/groovy/DataGen.groovy
@@ -31,10 +31,8 @@ class DataGen {
         )
     }
 
-    static Workspace workspace(String name, Map<String, String> agents = [:]) {
-        UUID wsId = UUID.randomUUID()
-        UUID ownerId = UUID.randomUUID()
-        List<WorkspaceResource> resources = agents.collect{ agentName, agentId -> new WorkspaceResource(
+    static WorkspaceResource resource(String agentName, String agentId){
+        return new WorkspaceResource(
                 UUID.randomUUID(),      // id
                 new Date().toInstant(), // created_at
                 UUID.randomUUID(),      // job_id
@@ -70,7 +68,13 @@ class DataGen {
                 )),
                 null,                   // metadata
                 0,                      // daily_cost
-        )}
+        )
+    }
+
+    static Workspace workspace(String name, Map<String, String> agents = [:]) {
+        UUID wsId = UUID.randomUUID()
+        UUID ownerId = UUID.randomUUID()
+        List<WorkspaceResource> resources = agents.collect{ resource(it.key, it.value)}
         return new Workspace(
                 wsId,
                 new Date().toInstant(), // created_at

From 4f61db7ab374399a422f11aa9cad99c16096deb9 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 17 Oct 2023 15:39:02 -0800
Subject: [PATCH 10/10] Add tests for workspaces and agents endpoints

---
 src/test/groovy/CoderRestClientTest.groovy | 119 ++++++++++++++++++++-
 1 file changed, 118 insertions(+), 1 deletion(-)

diff --git a/src/test/groovy/CoderRestClientTest.groovy b/src/test/groovy/CoderRestClientTest.groovy
index 3dd8dd3e..cadea39a 100644
--- a/src/test/groovy/CoderRestClientTest.groovy
+++ b/src/test/groovy/CoderRestClientTest.groovy
@@ -1,9 +1,126 @@
 package com.coder.gateway.sdk
 
-import spock.lang.*
+import com.coder.gateway.sdk.convertors.InstantConverter
+import com.coder.gateway.sdk.v2.models.Workspace
+import com.coder.gateway.sdk.v2.models.WorkspaceResource
+import com.coder.gateway.sdk.v2.models.WorkspacesResponse
+import com.google.gson.GsonBuilder
+import com.sun.net.httpserver.HttpExchange
+import com.sun.net.httpserver.HttpHandler
+import com.sun.net.httpserver.HttpServer
+import spock.lang.IgnoreIf
+import spock.lang.Requires
+import spock.lang.Specification
+import spock.lang.Unroll
+
+import javax.net.ssl.HttpsURLConnection
+import java.time.Instant
 
 @Unroll
 class CoderRestClientTest extends Specification {
+    /**
+     * Create, start, and return a server that mocks the Coder API.
+     *
+     * The resources map to the workspace index (to avoid having to manually hardcode IDs everywhere since you cannot
+     * use variables in the where blocks).
+     */
+    def mockServer(List<Workspace> workspaces, List<List<WorkspaceResource>> resources = []) {
+        HttpServer srv = HttpServer.create(new InetSocketAddress(0), 0)
+        srv.createContext("/", new HttpHandler() {
+            void handle(HttpExchange exchange) {
+                int code = HttpURLConnection.HTTP_NOT_FOUND
+                String response = "not found"
+                try {
+                    def matcher = exchange.requestURI.path =~ /\/api\/v2\/templateversions\/([^\/]+)\/resources/
+                    if (matcher.size() == 1) {
+                        UUID templateVersionId = UUID.fromString(matcher[0][1])
+                        int idx = workspaces.findIndexOf { it.latestBuild.templateVersionID == templateVersionId }
+                        code = HttpURLConnection.HTTP_OK
+                        response = new GsonBuilder().registerTypeAdapter(Instant.class, new InstantConverter())
+                                .create().toJson(resources[idx])
+                    } else if (exchange.requestURI.path == "/api/v2/workspaces") {
+                        code = HttpsURLConnection.HTTP_OK
+                        response = new GsonBuilder().registerTypeAdapter(Instant.class, new InstantConverter())
+                                .create().toJson(new WorkspacesResponse(workspaces, workspaces.size()))
+                    }
+                } catch (error) {
+                    // This will be a developer error.
+                    code = HttpURLConnection.HTTP_INTERNAL_ERROR
+                    response = error.message
+                    println(error.message) // Print since it will not show up in the error.
+                }
+
+                byte[] body = response.getBytes()
+                exchange.sendResponseHeaders(code, body.length)
+                exchange.responseBody.write(body)
+                exchange.close()
+            }
+        })
+        srv.start()
+        return [srv, "http://localhost:" + srv.address.port]
+    }
+
+    def "gets workspaces"() {
+        given:
+        def (srv, url) = mockServer(workspaces)
+        def client = new CoderRestClient(new URL(url), "token", null, "test")
+
+        expect:
+        client.workspaces()*.name == expected
+
+        cleanup:
+        srv.stop(0)
+
+        where:
+        workspaces                                           | expected
+        []                                                   | []
+        [DataGen.workspace("ws1")]                           | ["ws1"]
+        [DataGen.workspace("ws1"), DataGen.workspace("ws2")] | ["ws1", "ws2"]
+    }
+
+    def "gets resources"() {
+        given:
+        def (srv, url) = mockServer(workspaces, resources)
+        def client = new CoderRestClient(new URL(url), "token", null, "test")
+
+        expect:
+        client.agents(workspaces).collect { it.agentID.toString() } == expected
+
+        cleanup:
+        srv.stop(0)
+
+        where:
+        workspaces << [
+                [],
+                [DataGen.workspace("ws1", [agent1: "3f51da1d-306f-4a40-ac12-62bda5bc5f9a"])],
+                [DataGen.workspace("ws1", [agent1: "3f51da1d-306f-4a40-ac12-62bda5bc5f9a"])],
+                [DataGen.workspace("ws1", [agent1: "3f51da1d-306f-4a40-ac12-62bda5bc5f9a"]),
+                 DataGen.workspace("ws2"),
+                 DataGen.workspace("ws3")],
+        ]
+        resources << [
+                [],
+                [[]],
+                [[DataGen.resource("agent2", "968eea5e-8787-439d-88cd-5bc440216a34"),
+                  DataGen.resource("agent3", "72fbc97b-952c-40c8-b1e5-7535f4407728")]],
+                [[],
+                 [DataGen.resource("agent2", "968eea5e-8787-439d-88cd-5bc440216a34"),
+                  DataGen.resource("agent3", "72fbc97b-952c-40c8-b1e5-7535f4407728")],
+                 []],
+        ]
+        expected << [
+                // Nothing, so no agents.
+                [],
+                // One workspace with an agent, but resources get overridden by the resources endpoint that returns
+                // nothing so we end up with a workspace without an agent.
+                ["null"],
+                // One workspace with an agent, but resources get overridden by the resources endpoint.
+                ["968eea5e-8787-439d-88cd-5bc440216a34", "72fbc97b-952c-40c8-b1e5-7535f4407728"],
+                // Multiple workspaces but only one has resources from the resources endpoint.
+                ["null", "968eea5e-8787-439d-88cd-5bc440216a34", "72fbc97b-952c-40c8-b1e5-7535f4407728", "null"],
+        ]
+    }
+
     def "gets headers"() {
         expect:
         CoderRestClient.getHeaders(new URL("http://localhost"), command) == expected