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.

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)

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"
+    }
+}

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
+        )
+    }
 }