Skip to content

Commit 716a2fb

Browse files
mafredrimafredri
committed
add tests
1 parent 81d3f62 commit 716a2fb

File tree

1 file changed

+193
-0
lines changed

1 file changed

+193
-0
lines changed

‎cli/ssh_internal_test.go

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"time"
1212

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

0 commit comments

Comments
 (0)