chore: Write unit tests for agent selector

Fix fake database implementation for returning no rows
coder · Emyrk · Aug 29, 2022 · Aug 24, 2022 · Aug 24, 2022 · Aug 24, 2022
commit 0bb51b46ee4d99d212028adcbce7d47001897a1a
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -138,6 +138,7 @@ func New(options *Options) *API {
 			httpmw.RateLimitPerMinute(options.APIRateLimit),
 			httpmw.ExtractAPIKey(options.Database, oauthConfigs, true),
 			httpmw.ExtractUserParam(api.Database),
+			// Extracts the <workspace.agent> from the url
 			httpmw.ExtractWorkspaceAndAgentParam(api.Database),
 		)
 		r.HandleFunc("/*", api.workspaceAppsProxyPath)

diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -591,14 +591,14 @@ func (q *fakeQuerier) GetLatestWorkspaceBuildByWorkspaceID(_ context.Context, wo
 	defer q.mutex.RUnlock()
 
 	var row database.WorkspaceBuild
-	var buildNum int32
+	var buildNum int32 = -1
 	for _, workspaceBuild := range q.workspaceBuilds {
 		if workspaceBuild.WorkspaceID.String() == workspaceID.String() && workspaceBuild.BuildNumber > buildNum {
 			row = workspaceBuild
 			buildNum = workspaceBuild.BuildNumber
 		}
 	}
-	if buildNum == 0 {
+	if buildNum == -1 {
 		return database.WorkspaceBuild{}, sql.ErrNoRows
 	}
 	return row, nil
@@ -1262,9 +1262,6 @@ func (q *fakeQuerier) GetWorkspaceAgentsByResourceIDs(_ context.Context, resourc
 			workspaceAgents = append(workspaceAgents, agent)
 		}
 	}
-	if len(workspaceAgents) == 0 {
-		return nil, sql.ErrNoRows
-	}
 	return workspaceAgents, nil
 }
 
@@ -1346,9 +1343,6 @@ func (q *fakeQuerier) GetWorkspaceResourcesByJobID(_ context.Context, jobID uuid
 		}
 		resources = append(resources, resource)
 	}
-	if len(resources) == 0 {
-		return nil, sql.ErrNoRows
-	}
 	return resources, nil
 }
 

diff --git a/coderd/httpmw/workspaceparam.go b/coderd/httpmw/workspaceparam.go
@@ -149,8 +149,9 @@ func ExtractWorkspaceAndAgentParam(db database.Store) func(http.Handler) http.Ha
 				agent = agents[0]
 			}
 
-			ctx := context.WithValue(r.Context(), workspaceParamContextKey{}, workspace)
-			ctx = context.WithValue(r.Context(), workspaceAgentContextKey{}, agent)
+			ctx := r.Context()
+			ctx = context.WithValue(ctx, workspaceParamContextKey{}, workspace)
+			ctx = context.WithValue(ctx, workspaceAgentParamContextKey{}, agent)
 			next.ServeHTTP(rw, r.WithContext(ctx))
 		})
 	}

diff --git a/coderd/httpmw/workspaceparam_test.go b/coderd/httpmw/workspaceparam_test.go
@@ -3,12 +3,15 @@ package httpmw_test
 import (
 	"context"
 	"crypto/sha256"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
@@ -122,3 +125,240 @@ func TestWorkspaceParam(t *testing.T) {
 		require.Equal(t, http.StatusOK, res.StatusCode)
 	})
 }
+
+func TestWorkspaceAgentByNameParam(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		Name string
+		// Agents are mapped to a resource
+		Agents             map[string][]string
+		UrlParam           string
+		WorkspaceName      string
+		ExpectedAgent      string
+		ExpectedStatusCode int
+		ExpectedError      string
+	}{
+		{
+			Name:               "NoAgents",
+			WorkspaceName:      "dev",
+			Agents:             map[string][]string{},
+			UrlParam:           "dev",
+			ExpectedError:      "No agents exist",
+			ExpectedStatusCode: http.StatusBadRequest,
+		},
+		{
+			Name:          "MultipleAgents",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {
+					"agent-one",
+					"agent-two",
+				},
+			},
+			UrlParam:           "dev",
+			ExpectedStatusCode: http.StatusBadRequest,
+			ExpectedError:      "More than one agent exists, but no agent specified",
+		},
+		{
+			Name:          "MultipleResources",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {
+					"agent-one",
+				},
+				"resource-b": {
+					"agent-two",
+				},
+			},
+			UrlParam:           "dev",
+			ExpectedStatusCode: http.StatusBadRequest,
+			ExpectedError:      "More than one agent exists, but no agent specified",
+		},
+
+		// OKs
+		{
+			Name:          "MultipleResourcesOneAgent",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {},
+				"resource-b": {
+					"agent-one",
+				},
+			},
+			UrlParam:           "dev",
+			ExpectedAgent:      "agent-one",
+			ExpectedStatusCode: http.StatusOK,
+		},
+		{
+			Name:          "OneAgent",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {
+					"agent-one",
+				},
+			},
+			UrlParam:           "dev",
+			ExpectedAgent:      "agent-one",
+			ExpectedStatusCode: http.StatusOK,
+		},
+		{
+			Name:          "MultipleAgentSelectOne",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {
+					"agent-one",
+					"agent-two",
+					"agent-selected",
+				},
+			},
+			UrlParam:           "dev.agent-selected",
+			ExpectedAgent:      "agent-selected",
+			ExpectedStatusCode: http.StatusOK,
+		},
+		{
+			Name:          "MultipleResourcesSelectOne",
+			WorkspaceName: "dev",
+			Agents: map[string][]string{
+				"resource-a": {
+					"agent-one",
+				},
+				"resource-b": {
+					"agent-two",
+				},
+				"resource-c": {
+					"agent-selected",
+					"agent-three",
+				},
+			},
+			UrlParam:           "dev.agent-selected",
+			ExpectedAgent:      "agent-selected",
+			ExpectedStatusCode: http.StatusOK,
+		},
+	}
+
+	for _, c := range testCases {
+		c := c
+		t.Run(c.Name, func(t *testing.T) {
+			db, r := setupWorkspaceWithAgents(t, setupConfig{
+				WorkspaceName: c.WorkspaceName,
+				Agents:        c.Agents,
+			})
+
+			chi.RouteContext(r.Context()).URLParams.Add("workspacename_and_agent", c.UrlParam)
+
+			rtr := chi.NewRouter()
+			rtr.Use(
+				httpmw.ExtractAPIKey(db, nil, true),
+				httpmw.ExtractUserParam(db),
+				httpmw.ExtractWorkspaceAndAgentParam(db),
+			)
+			rtr.Get("/", func(w http.ResponseWriter, r *http.Request) {
+				workspace := httpmw.WorkspaceParam(r)
+				agent := httpmw.WorkspaceAgentParam(r)
+
+				assert.Equal(t, c.ExpectedAgent, agent.Name, "expected agent name")
+				assert.Equal(t, c.WorkspaceName, workspace.Name, "expected workspace name")
+			})
+
+			rw := httptest.NewRecorder()
+			rtr.ServeHTTP(rw, r)
+			res := rw.Result()
+			var coderResp codersdk.Response
+			_ = json.NewDecoder(res.Body).Decode(&coderResp)
+			res.Body.Close()
+			require.Equal(t, c.ExpectedStatusCode, res.StatusCode)
+			if c.ExpectedError != "" {
+				require.Contains(t, coderResp.Message, c.ExpectedError)
+			}
+		})
+	}
+}
+
+type setupConfig struct {
+	WorkspaceName string
+	// Agents are mapped to a resource
+	Agents map[string][]string
+}
+
+func setupWorkspaceWithAgents(t testing.TB, cfg setupConfig) (database.Store, *http.Request) {
+	t.Helper()
+	db := databasefake.New()
+	var (
+		id, secret = randomAPIKeyParts()
+		hashed     = sha256.Sum256([]byte(secret))
+	)
+	r := httptest.NewRequest("GET", "/", nil)
+	r.AddCookie(&http.Cookie{
+		Name:  codersdk.SessionTokenKey,
+		Value: fmt.Sprintf("%s-%s", id, secret),
+	})
+
+	userID := uuid.New()
+	username, err := cryptorand.String(8)
+	require.NoError(t, err)
+	user, err := db.InsertUser(r.Context(), database.InsertUserParams{
+		ID:             userID,
+		Email:          "testaccount@coder.com",
+		HashedPassword: hashed[:],
+		Username:       username,
+		CreatedAt:      database.Now(),
+		UpdatedAt:      database.Now(),
+	})
+	require.NoError(t, err)
+
+	_, err = db.InsertAPIKey(r.Context(), database.InsertAPIKeyParams{
+		ID:           id,
+		UserID:       user.ID,
+		HashedSecret: hashed[:],
+		LastUsed:     database.Now(),
+		ExpiresAt:    database.Now().Add(time.Minute),
+		LoginType:    database.LoginTypePassword,
+	})
+	require.NoError(t, err)
+
+	workspace, err := db.InsertWorkspace(context.Background(), database.InsertWorkspaceParams{
+		ID:         uuid.New(),
+		TemplateID: uuid.New(),
+		OwnerID:    user.ID,
+		Name:       cfg.WorkspaceName,
+	})
+	require.NoError(t, err)
+
+	build, err := db.InsertWorkspaceBuild(context.Background(), database.InsertWorkspaceBuildParams{
+		ID:          uuid.New(),
+		WorkspaceID: workspace.ID,
+		JobID:       uuid.New(),
+	})
+	require.NoError(t, err)
+
+	job, err := db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
+		ID:   build.JobID,
+		Type: database.ProvisionerJobTypeWorkspaceBuild,
+	})
+	require.NoError(t, err)
+
+	for resourceName, agentNames := range cfg.Agents {
+		resource, err := db.InsertWorkspaceResource(context.Background(), database.InsertWorkspaceResourceParams{
+			ID:    uuid.New(),
+			JobID: job.ID,
+			Name:  resourceName,
+		})
+		require.NoError(t, err)
+
+		for _, name := range agentNames {
+			_, err = db.InsertWorkspaceAgent(context.Background(), database.InsertWorkspaceAgentParams{
+				ID:         uuid.New(),
+				ResourceID: resource.ID,
+				Name:       name,
+			})
+			require.NoError(t, err)
+		}
+	}
+
+	ctx := chi.NewRouteContext()
+	ctx.URLParams.Add("user", codersdk.Me)
+	r = r.WithContext(context.WithValue(r.Context(), chi.RouteCtxKey, ctx))
+
+	return db, r
+}