Skip to content

Commit eef5bb3

Browse files
committed
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.
1 parent 4f26d38 commit eef5bb3

File tree

3 files changed

+253
-26
lines changed

3 files changed

+253
-26
lines changed

src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt

+48-24
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
package com.coder.gateway
44

55
import com.coder.gateway.models.TokenSource
6+
import com.coder.gateway.models.WorkspaceAgentModel
67
import com.coder.gateway.sdk.CoderCLIManager
78
import com.coder.gateway.sdk.CoderRestClient
89
import com.coder.gateway.sdk.ex.AuthenticationResponseException
910
import com.coder.gateway.sdk.toURL
11+
import com.coder.gateway.sdk.v2.models.Workspace
1012
import com.coder.gateway.sdk.v2.models.WorkspaceStatus
1113
import com.coder.gateway.sdk.v2.models.toAgentModels
1214
import com.coder.gateway.sdk.withPath
@@ -67,30 +69,8 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
6769
WorkspaceStatus.RUNNING -> Unit // All is well
6870
}
6971

70-
val agents = workspace.toAgentModels()
71-
if (agents.isEmpty()) {
72-
throw IllegalArgumentException("The workspace \"$workspaceName\" has no agents")
73-
}
74-
75-
// If the agent is missing and the workspace has only one, use that.
76-
// Prefer the ID over the name if both are set.
77-
val agent = if (!parameters[AGENT_ID].isNullOrBlank())
78-
agents.firstOrNull {it.agentID.toString() == parameters[AGENT_ID]}
79-
else if (!parameters[AGENT_NAME].isNullOrBlank())
80-
agents.firstOrNull { it.name == "$workspaceName.${parameters[AGENT_NAME]}"}
81-
else if (agents.size == 1) agents.first()
82-
else null
83-
84-
if (agent == null) {
85-
if (parameters[AGENT_ID].isNullOrBlank() && parameters[AGENT_NAME].isNullOrBlank()) {
86-
// TODO: Show a dropdown and ask for an agent.
87-
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")
88-
} else if (parameters[AGENT_ID].isNullOrBlank()) {
89-
throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
90-
} else {
91-
throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent named \"${parameters[AGENT_NAME]}\"")
92-
}
93-
}
72+
// TODO: Show a dropdown and ask for an agent if missing.
73+
val agent = getMatchingAgent(parameters, workspace)
9474

9575
if (agent.agentStatus.pending()) {
9676
// TODO: Wait for the agent to be ready.
@@ -211,5 +191,49 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
211191

212192
companion object {
213193
val logger = Logger.getInstance(CoderGatewayConnectionProvider::class.java.simpleName)
194+
195+
/**
196+
* Return the agent matching the provided agent ID or name in the
197+
* parameters. The name is ignored if the ID is set. If neither was
198+
* supplied and the workspace has only one agent, return that.
199+
* Otherwise throw an error.
200+
*
201+
* @throws [MissingArgumentException, IllegalArgumentException]
202+
*/
203+
@JvmStatic
204+
fun getMatchingAgent(parameters: Map<String, String>, workspace: Workspace): WorkspaceAgentModel {
205+
// A WorkspaceAgentModel will still be returned if there are no
206+
// agents; in this case it represents the workspace instead.
207+
// TODO: Seems confusing for something with "agent" in the name to
208+
// potentially not actually be an agent; can we replace
209+
// WorkspaceAgentModel with the original structs from the API?
210+
val agents = workspace.toAgentModels()
211+
if (agents.isEmpty() || (agents.size == 1 && agents.first().agentID == null)) {
212+
throw IllegalArgumentException("The workspace \"${workspace.name}\" has no agents")
213+
}
214+
215+
// If the agent is missing and the workspace has only one, use that.
216+
// Prefer the ID over the name if both are set.
217+
val agent = if (!parameters[AGENT_ID].isNullOrBlank())
218+
agents.firstOrNull { it.agentID.toString() == parameters[AGENT_ID] }
219+
else if (!parameters[AGENT_NAME].isNullOrBlank())
220+
agents.firstOrNull { it.name == "${workspace.name}.${parameters[AGENT_NAME]}"}
221+
else if (agents.size == 1) agents.first()
222+
else null
223+
224+
if (agent == null) {
225+
if (!parameters[AGENT_ID].isNullOrBlank()) {
226+
throw IllegalArgumentException("The workspace \"${workspace.name}\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
227+
} else if (!parameters[AGENT_NAME].isNullOrBlank()){
228+
throw IllegalArgumentException("The workspace \"${workspace.name}\"does not have an agent named \"${parameters[AGENT_NAME]}\"")
229+
} else {
230+
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")
231+
}
232+
}
233+
234+
return agent
235+
}
214236
}
215237
}
238+
239+
class MissingArgumentException(message: String) : IllegalArgumentException(message)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package com.coder.gateway
2+
3+
import spock.lang.Shared
4+
import spock.lang.Specification
5+
import spock.lang.Unroll
6+
7+
@Unroll
8+
class CoderGatewayConnectionProviderTest extends Specification {
9+
@Shared
10+
def agents = [
11+
agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24",
12+
agent_name_2: "fb3daea4-da6b-424d-84c7-36b90574cfef",
13+
agent_name: "9a920eee-47fb-4571-9501-e4b3120c12f2",
14+
]
15+
def oneAgent = [
16+
agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
17+
]
18+
19+
def "gets matching agent"() {
20+
expect:
21+
def ws = DataGen.workspace("ws", agents)
22+
CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString(expected)
23+
24+
where:
25+
parameters | expected
26+
[agent: "agent_name"] | "9a920eee-47fb-4571-9501-e4b3120c12f2"
27+
[agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | "9a920eee-47fb-4571-9501-e4b3120c12f2"
28+
[agent: "agent_name_2"] | "fb3daea4-da6b-424d-84c7-36b90574cfef"
29+
[agent_id: "fb3daea4-da6b-424d-84c7-36b90574cfef"] | "fb3daea4-da6b-424d-84c7-36b90574cfef"
30+
[agent: "agent_name_3"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
31+
[agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
32+
33+
// Prefer agent_id.
34+
[agent: "agent_name", agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
35+
}
36+
37+
def "fails to get matching agent"() {
38+
when:
39+
def ws = DataGen.workspace("ws", agents)
40+
CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
41+
42+
then:
43+
def err = thrown(expected)
44+
err.message.contains(message)
45+
46+
where:
47+
parameters | expected | message
48+
[:] | MissingArgumentException | "Unable to determine"
49+
[agent: ""] | MissingArgumentException | "Unable to determine"
50+
[agent_id: ""] | MissingArgumentException | "Unable to determine"
51+
[agent: null] | MissingArgumentException | "Unable to determine"
52+
[agent_id: null] | MissingArgumentException | "Unable to determine"
53+
[agent: "ws"] | IllegalArgumentException | "agent named"
54+
[agent: "ws.agent_name"] | IllegalArgumentException | "agent named"
55+
[agent: "agent_name_4"] | IllegalArgumentException | "agent named"
56+
[agent_id: "not-a-uuid"] | IllegalArgumentException | "agent with ID"
57+
[agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
58+
59+
// Will ignore agent if agent_id is set even if agent matches.
60+
[agent: "agent_name", agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
61+
}
62+
63+
def "gets the first agent when workspace has only one"() {
64+
expect:
65+
def ws = DataGen.workspace("ws", oneAgent)
66+
CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString("b0e4c54d-9ba9-4413-8512-11ca1e826a24")
67+
68+
where:
69+
parameters << [
70+
[:],
71+
[agent: ""],
72+
[agent_id: ""],
73+
[agent: null],
74+
[agent_id: null],
75+
]
76+
}
77+
78+
def "fails to get agent when workspace has only one"() {
79+
when:
80+
def ws = DataGen.workspace("ws", oneAgent)
81+
CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
82+
83+
then:
84+
def err = thrown(expected)
85+
err.message.contains(message)
86+
87+
where:
88+
parameters | expected | message
89+
[agent: "ws"] | IllegalArgumentException | "agent named"
90+
[agent: "ws.agent_name_3"] | IllegalArgumentException | "agent named"
91+
[agent: "agent_name_4"] | IllegalArgumentException | "agent named"
92+
[agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
93+
}
94+
95+
def "fails to get agent from workspace without agents"() {
96+
when:
97+
def ws = DataGen.workspace("ws")
98+
CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
99+
100+
then:
101+
def err = thrown(expected)
102+
err.message.contains(message)
103+
104+
where:
105+
parameters | expected | message
106+
[:] | IllegalArgumentException | "has no agents"
107+
[agent: ""] | IllegalArgumentException | "has no agents"
108+
[agent_id: ""] | IllegalArgumentException | "has no agents"
109+
[agent: null] | IllegalArgumentException | "has no agents"
110+
[agent_id: null] | IllegalArgumentException | "has no agents"
111+
[agent: "agent_name"] | IllegalArgumentException | "has no agents"
112+
[agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | IllegalArgumentException | "has no agents"
113+
}
114+
}

src/test/groovy/DataGen.groovy

+91-2
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
import com.coder.gateway.models.WorkspaceAgentModel
22
import com.coder.gateway.models.WorkspaceAndAgentStatus
33
import com.coder.gateway.models.WorkspaceVersionStatus
4-
import com.coder.gateway.sdk.v2.models.WorkspaceStatus
5-
import com.coder.gateway.sdk.v2.models.WorkspaceTransition
4+
import com.coder.gateway.sdk.v2.models.*
65

76
class DataGen {
87
// Create a random workspace agent model. If the workspace name is omitted
@@ -31,4 +30,94 @@ class DataGen {
3130
null
3231
)
3332
}
33+
34+
static Workspace workspace(String name, Map<String, String> agents = [:]) {
35+
UUID wsId = UUID.randomUUID()
36+
UUID ownerId = UUID.randomUUID()
37+
List<WorkspaceResource> resources = agents.collect{ agentName, agentId -> new WorkspaceResource(
38+
UUID.randomUUID(), // id
39+
new Date().toInstant(), // created_at
40+
UUID.randomUUID(), // job_id
41+
WorkspaceTransition.START,
42+
"type",
43+
"name",
44+
false, // hide
45+
"icon",
46+
List.of(new WorkspaceAgent(
47+
UUID.fromString(agentId),
48+
new Date().toInstant(), // created_at
49+
new Date().toInstant(), // updated_at
50+
null, // first_connected_at
51+
null, // last_connected_at
52+
null, // disconnected_at
53+
WorkspaceAgentStatus.CONNECTED,
54+
agentName,
55+
UUID.randomUUID(), // resource_id
56+
null, // instance_id
57+
"arch", // architecture
58+
[:], // environment_variables
59+
"os", // operating_system
60+
null, // startup_script
61+
null, // directory
62+
null, // expanded_directory
63+
"version", // version
64+
List.of(), // apps
65+
null, // latency
66+
0, // connection_timeout_seconds
67+
"url", // troubleshooting_url
68+
WorkspaceAgentLifecycleState.READY,
69+
false, // login_before_ready
70+
)),
71+
null, // metadata
72+
0, // daily_cost
73+
)}
74+
return new Workspace(
75+
wsId,
76+
new Date().toInstant(), // created_at
77+
new Date().toInstant(), // updated_at
78+
ownerId,
79+
"owner-name",
80+
UUID.randomUUID(), // template_id
81+
"template-name",
82+
"template-display-name",
83+
"template-icon",
84+
false, // template_allow_user_cancel_workspace_jobs
85+
new WorkspaceBuild(
86+
UUID.randomUUID(), // id
87+
new Date().toInstant(), // created_at
88+
new Date().toInstant(), // updated_at
89+
wsId,
90+
name,
91+
ownerId,
92+
"owner-name",
93+
UUID.randomUUID(), // template_version_id
94+
0, // build_number
95+
WorkspaceTransition.START,
96+
UUID.randomUUID(), // initiator_id
97+
"initiator-name",
98+
new ProvisionerJob(
99+
UUID.randomUUID(), // id
100+
new Date().toInstant(), // created_at
101+
null, // started_at
102+
null, // completed_at
103+
null, // canceled_at
104+
null, // error
105+
ProvisionerJobStatus.SUCCEEDED,
106+
null, // worker_id
107+
UUID.randomUUID(), // file_id
108+
[:], // tags
109+
),
110+
BuildReason.INITIATOR,
111+
resources,
112+
null, // deadline
113+
WorkspaceStatus.RUNNING,
114+
0, // daily_cost
115+
),
116+
false, // outdated
117+
name,
118+
null, // autostart_schedule
119+
null, // ttl_ms
120+
new Date().toInstant(), // last_used_at
121+
)
122+
}
34123
}

0 commit comments

Comments
 (0)