Skip to content

Commit 1c6193a

Browse files
mafredrimafredri
committed
add tests
1 parent 647181e commit 1c6193a

File tree

1 file changed

+192
-0
lines changed

1 file changed

+192
-0
lines changed

‎cli/ssh_internal_test.go

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,3 +346,195 @@ func newAsyncCloser(ctx context.Context, t *testing.T) *asyncCloser {
346346
started: make(chan struct{}),
347347
}
348348
}
349+
350+
func Test_getWorkspaceAgent(t *testing.T) {
351+
t.Parallel()
352+
353+
createWorkspaceWithAgents := func(agents []codersdk.WorkspaceAgent) codersdk.Workspace {
354+
return codersdk.Workspace{
355+
Name: "test-workspace",
356+
LatestBuild: codersdk.WorkspaceBuild{
357+
Resources: []codersdk.WorkspaceResource{
358+
{
359+
Agents: agents,
360+
},
361+
},
362+
},
363+
}
364+
}
365+
366+
createAgent := func(name string) codersdk.WorkspaceAgent {
367+
return codersdk.WorkspaceAgent{
368+
ID: uuid.New(),
369+
Name: name,
370+
ParentID: uuid.NullUUID{},
371+
}
372+
}
373+
374+
createSubAgent := func(name string, parentID uuid.UUID) codersdk.WorkspaceAgent {
375+
return codersdk.WorkspaceAgent{
376+
ID: uuid.New(),
377+
Name: name,
378+
ParentID: uuid.NullUUID{
379+
UUID: parentID,
380+
Valid: true,
381+
},
382+
}
383+
}
384+
385+
t.Run("SingleAgent_NoNameSpecified", func(t *testing.T) {
386+
t.Parallel()
387+
agent := createAgent("main")
388+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent})
389+
390+
result, err := getWorkspaceAgent(workspace, "")
391+
require.NoError(t, err)
392+
assert.Equal(t, agent.ID, result.ID)
393+
assert.Equal(t, "main", result.Name)
394+
})
395+
396+
t.Run("SingleSubAgent_NoNameSpecified", func(t *testing.T) {
397+
t.Parallel()
398+
parentAgent := createAgent("main")
399+
subAgent := createSubAgent("devcontainer", parentAgent.ID)
400+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent})
401+
402+
// Should prefer the sub-agent when no name is specified.
403+
result, err := getWorkspaceAgent(workspace, "")
404+
require.NoError(t, err)
405+
assert.Equal(t, subAgent.ID, result.ID)
406+
assert.Equal(t, "devcontainer", result.Name)
407+
})
408+
409+
t.Run("MultipleAgents_NoSubAgents_NoNameSpecified", func(t *testing.T) {
410+
t.Parallel()
411+
agent1 := createAgent("main1")
412+
agent2 := createAgent("main2")
413+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
414+
415+
_, err := getWorkspaceAgent(workspace, "")
416+
require.Error(t, err)
417+
assert.Contains(t, err.Error(), "multiple agents found")
418+
assert.Contains(t, err.Error(), "available agents: [main1 main2]")
419+
})
420+
421+
t.Run("MultipleSubAgents_NoNameSpecified", func(t *testing.T) {
422+
t.Parallel()
423+
parentAgent := createAgent("main")
424+
subAgent1 := createSubAgent("devcontainer1", parentAgent.ID)
425+
subAgent2 := createSubAgent("devcontainer2", parentAgent.ID)
426+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent1, subAgent2})
427+
428+
_, err := getWorkspaceAgent(workspace, "")
429+
require.Error(t, err)
430+
assert.Contains(t, err.Error(), "multiple sub-agents found")
431+
assert.Contains(t, err.Error(), "available agents: [devcontainer1 devcontainer2 main]")
432+
})
433+
434+
t.Run("AgentNameSpecified_Found_RegularAgent", func(t *testing.T) {
435+
t.Parallel()
436+
agent1 := createAgent("main1")
437+
agent2 := createAgent("main2")
438+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
439+
440+
result, err := getWorkspaceAgent(workspace, "main1")
441+
require.NoError(t, err)
442+
assert.Equal(t, agent1.ID, result.ID)
443+
assert.Equal(t, "main1", result.Name)
444+
})
445+
446+
t.Run("AgentNameSpecified_Found_SubAgent", func(t *testing.T) {
447+
t.Parallel()
448+
agent := createAgent("main")
449+
subAgent := createSubAgent("devcontainer", agent.ID)
450+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
451+
452+
result, err := getWorkspaceAgent(workspace, "devcontainer")
453+
require.NoError(t, err)
454+
assert.Equal(t, subAgent.ID, result.ID)
455+
assert.Equal(t, "devcontainer", result.Name)
456+
})
457+
458+
t.Run("AgentNameSpecified_NotFound", func(t *testing.T) {
459+
t.Parallel()
460+
agent1 := createAgent("main1")
461+
agent2 := createAgent("main2")
462+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
463+
464+
_, err := getWorkspaceAgent(workspace, "nonexistent")
465+
require.Error(t, err)
466+
assert.Contains(t, err.Error(), `agent not found by name "nonexistent"`)
467+
assert.Contains(t, err.Error(), "available agents: [main1 main2]")
468+
})
469+
470+
t.Run("NoAgents", func(t *testing.T) {
471+
t.Parallel()
472+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{})
473+
474+
_, err := getWorkspaceAgent(workspace, "")
475+
require.Error(t, err)
476+
assert.Contains(t, err.Error(), `workspace "test-workspace" has no agents`)
477+
})
478+
479+
t.Run("MixedAgents_SubAgentPreferred", func(t *testing.T) {
480+
t.Parallel()
481+
agent := createAgent("main")
482+
subAgent := createSubAgent("devcontainer", agent.ID)
483+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
484+
485+
// When no name is specified and there's one sub-agent,
486+
// it should be preferred.
487+
result, err := getWorkspaceAgent(workspace, "")
488+
require.NoError(t, err)
489+
assert.Equal(t, subAgent.ID, result.ID)
490+
assert.Equal(t, "devcontainer", result.Name)
491+
})
492+
493+
t.Run("MixedAgents_SpecificNameFound", func(t *testing.T) {
494+
t.Parallel()
495+
agent := createAgent("main")
496+
subAgent := createSubAgent("devcontainer", agent.ID)
497+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
498+
499+
// Should be able to find regular agent by name.
500+
result, err := getWorkspaceAgent(workspace, "main")
501+
require.NoError(t, err)
502+
assert.Equal(t, agent.ID, result.ID)
503+
assert.Equal(t, "main", result.Name)
504+
505+
// Should be able to find sub-agent by name.
506+
result, err = getWorkspaceAgent(workspace, "devcontainer")
507+
require.NoError(t, err)
508+
assert.Equal(t, subAgent.ID, result.ID)
509+
assert.Equal(t, "devcontainer", result.Name)
510+
})
511+
512+
t.Run("AvailableAgentNames_SortedCorrectly", func(t *testing.T) {
513+
t.Parallel()
514+
// Define agents in non-alphabetical order.
515+
agent2 := createAgent("zod")
516+
agent1 := createAgent("clark")
517+
subAgent := createSubAgent("krypton", agent1.ID)
518+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent2, agent1, subAgent})
519+
520+
_, err := getWorkspaceAgent(workspace, "nonexistent")
521+
require.Error(t, err)
522+
// Available agents should be sorted alphabetically.
523+
assert.Contains(t, err.Error(), "available agents: [clark krypton zod]")
524+
})
525+
526+
t.Run("MultipleAgentsAndSubAgents_NoNameSpecified", func(t *testing.T) {
527+
t.Parallel()
528+
agent1 := createAgent("main1")
529+
agent2 := createAgent("main2")
530+
subAgent1 := createSubAgent("dev1", agent1.ID)
531+
subAgent2 := createSubAgent("dev2", agent1.ID)
532+
workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2, subAgent1, subAgent2})
533+
534+
// Should error because there are multiple sub-agents.
535+
_, err := getWorkspaceAgent(workspace, "")
536+
require.Error(t, err)
537+
assert.Contains(t, err.Error(), "multiple sub-agents found")
538+
assert.Contains(t, err.Error(), "available agents: [dev1 dev2 main1 main2]")
539+
})
540+
}

0 commit comments

Comments
 (0)