add tests

mafredri · mafredri · commit f923edb6cb44 · 2025-06-18T15:05:58.000Z
diff --git a/cli/ssh_internal_test.go b/cli/ssh_internal_test.go
@@ -11,6 +11,7 @@ import (
 	"time"
 
 	gliderssh "github.com/gliderlabs/ssh"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/crypto/ssh"
@@ -346,3 +347,195 @@ func newAsyncCloser(ctx context.Context, t *testing.T) *asyncCloser {
 		started:    make(chan struct{}),
 	}
 }
+
+func Test_getWorkspaceAgent(t *testing.T) {
+	t.Parallel()
+
+	createWorkspaceWithAgents := func(agents []codersdk.WorkspaceAgent) codersdk.Workspace {
+		return codersdk.Workspace{
+			Name: "test-workspace",
+			LatestBuild: codersdk.WorkspaceBuild{
+				Resources: []codersdk.WorkspaceResource{
+					{
+						Agents: agents,
+					},
+				},
+			},
+		}
+	}
+
+	createAgent := func(name string) codersdk.WorkspaceAgent {
+		return codersdk.WorkspaceAgent{
+			ID:       uuid.New(),
+			Name:     name,
+			ParentID: uuid.NullUUID{},
+		}
+	}
+
+	createSubAgent := func(name string, parentID uuid.UUID) codersdk.WorkspaceAgent {
+		return codersdk.WorkspaceAgent{
+			ID:   uuid.New(),
+			Name: name,
+			ParentID: uuid.NullUUID{
+				UUID:  parentID,
+				Valid: true,
+			},
+		}
+	}
+
+	t.Run("SingleAgent_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		agent := createAgent("main")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent})
+
+		result, err := getWorkspaceAgent(workspace, "")
+		require.NoError(t, err)
+		assert.Equal(t, agent.ID, result.ID)
+		assert.Equal(t, "main", result.Name)
+	})
+
+	t.Run("SingleSubAgent_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		parentAgent := createAgent("main")
+		subAgent := createSubAgent("devcontainer", parentAgent.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent})
+
+		// Should prefer the sub-agent when no name is specified.
+		result, err := getWorkspaceAgent(workspace, "")
+		require.NoError(t, err)
+		assert.Equal(t, subAgent.ID, result.ID)
+		assert.Equal(t, "devcontainer", result.Name)
+	})
+
+	t.Run("MultipleAgents_NoSubAgents_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "multiple agents found")
+		assert.Contains(t, err.Error(), "available agents: [main1 main2]")
+	})
+
+	t.Run("MultipleSubAgents_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		parentAgent := createAgent("main")
+		subAgent1 := createSubAgent("devcontainer1", parentAgent.ID)
+		subAgent2 := createSubAgent("devcontainer2", parentAgent.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{parentAgent, subAgent1, subAgent2})
+
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "multiple sub-agents found")
+		assert.Contains(t, err.Error(), "available agents: [devcontainer1 devcontainer2 main]")
+	})
+
+	t.Run("AgentNameSpecified_Found_RegularAgent", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		result, err := getWorkspaceAgent(workspace, "main1")
+		require.NoError(t, err)
+		assert.Equal(t, agent1.ID, result.ID)
+		assert.Equal(t, "main1", result.Name)
+	})
+
+	t.Run("AgentNameSpecified_Found_SubAgent", func(t *testing.T) {
+		t.Parallel()
+		agent := createAgent("main")
+		subAgent := createSubAgent("devcontainer", agent.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
+
+		result, err := getWorkspaceAgent(workspace, "devcontainer")
+		require.NoError(t, err)
+		assert.Equal(t, subAgent.ID, result.ID)
+		assert.Equal(t, "devcontainer", result.Name)
+	})
+
+	t.Run("AgentNameSpecified_NotFound", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2})
+
+		_, err := getWorkspaceAgent(workspace, "nonexistent")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), `agent not found by name "nonexistent"`)
+		assert.Contains(t, err.Error(), "available agents: [main1 main2]")
+	})
+
+	t.Run("NoAgents", func(t *testing.T) {
+		t.Parallel()
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{})
+
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), `workspace "test-workspace" has no agents`)
+	})
+
+	t.Run("MixedAgents_SubAgentPreferred", func(t *testing.T) {
+		t.Parallel()
+		agent := createAgent("main")
+		subAgent := createSubAgent("devcontainer", agent.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
+
+		// When no name is specified and there's one sub-agent,
+		// it should be preferred.
+		result, err := getWorkspaceAgent(workspace, "")
+		require.NoError(t, err)
+		assert.Equal(t, subAgent.ID, result.ID)
+		assert.Equal(t, "devcontainer", result.Name)
+	})
+
+	t.Run("MixedAgents_SpecificNameFound", func(t *testing.T) {
+		t.Parallel()
+		agent := createAgent("main")
+		subAgent := createSubAgent("devcontainer", agent.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent, subAgent})
+
+		// Should be able to find regular agent by name.
+		result, err := getWorkspaceAgent(workspace, "main")
+		require.NoError(t, err)
+		assert.Equal(t, agent.ID, result.ID)
+		assert.Equal(t, "main", result.Name)
+
+		// Should be able to find sub-agent by name.
+		result, err = getWorkspaceAgent(workspace, "devcontainer")
+		require.NoError(t, err)
+		assert.Equal(t, subAgent.ID, result.ID)
+		assert.Equal(t, "devcontainer", result.Name)
+	})
+
+	t.Run("AvailableAgentNames_SortedCorrectly", func(t *testing.T) {
+		t.Parallel()
+		// Define agents in non-alphabetical order.
+		agent2 := createAgent("zod")
+		agent1 := createAgent("clark")
+		subAgent := createSubAgent("krypton", agent1.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent2, agent1, subAgent})
+
+		_, err := getWorkspaceAgent(workspace, "nonexistent")
+		require.Error(t, err)
+		// Available agents should be sorted alphabetically.
+		assert.Contains(t, err.Error(), "available agents: [clark krypton zod]")
+	})
+
+	t.Run("MultipleAgentsAndSubAgents_NoNameSpecified", func(t *testing.T) {
+		t.Parallel()
+		agent1 := createAgent("main1")
+		agent2 := createAgent("main2")
+		subAgent1 := createSubAgent("dev1", agent1.ID)
+		subAgent2 := createSubAgent("dev2", agent1.ID)
+		workspace := createWorkspaceWithAgents([]codersdk.WorkspaceAgent{agent1, agent2, subAgent1, subAgent2})
+
+		// Should error because there are multiple sub-agents.
+		_, err := getWorkspaceAgent(workspace, "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "multiple sub-agents found")
+		assert.Contains(t, err.Error(), "available agents: [dev1 dev2 main1 main2]")
+	})
+}
