coder/coderd/httpmw/workspaceagent_test.go

package httpmw_test

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	"github.com/stretchr/testify/require"

	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/database/dbgen"
	"github.com/coder/coder/v2/coderd/database/dbtestutil"
	"github.com/coder/coder/v2/coderd/httpmw"
	"github.com/coder/coder/v2/codersdk"
)

func TestWorkspaceAgent(t *testing.T) {
	t.Parallel()

	t.Run("None", func(t *testing.T) {
		t.Parallel()
		db, _ := dbtestutil.NewDB(t)

		req, rtr, _, _ := setup(t, db, uuid.New(), httpmw.ExtractWorkspaceAgentAndLatestBuild(
			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
				DB:       db,
				Optional: false,
			}))

		rw := httptest.NewRecorder()
		req.Header.Set(codersdk.SessionTokenHeader, uuid.New().String())
		rtr.ServeHTTP(rw, req)

		res := rw.Result()
		defer res.Body.Close()
		require.Equal(t, http.StatusUnauthorized, res.StatusCode)
	})

	t.Run("Found", func(t *testing.T) {
		t.Parallel()
		db, _ := dbtestutil.NewDB(t)
		authToken := uuid.New()
		req, rtr, _, _ := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
				DB:       db,
				Optional: false,
			}))

		rw := httptest.NewRecorder()
		req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
		rtr.ServeHTTP(rw, req)

		//nolint:bodyclose // Closed in `t.Cleanup`
		res := rw.Result()
		t.Cleanup(func() { _ = res.Body.Close() })
		require.Equal(t, http.StatusOK, res.StatusCode)
	})

	t.Run("Latest", func(t *testing.T) {
		t.Parallel()
		db, _ := dbtestutil.NewDB(t)
		authToken := uuid.New()
		req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
				DB:       db,
				Optional: false,
			}),
		)

		// Create a newer build
		job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
			OrganizationID: ws.OrganizationID,
		})
		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
			JobID: job.ID,
		})
		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
			WorkspaceID:       ws.ID,
			JobID:             job.ID,
			TemplateVersionID: tpv.ID,
			BuildNumber:       2,
		})
		_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
			ResourceID: resource.ID,
		})

		rw := httptest.NewRecorder()
		req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
		rtr.ServeHTTP(rw, req)

		//nolint:bodyclose // Closed in `t.Cleanup`
		res := rw.Result()
		t.Cleanup(func() { _ = res.Body.Close() })
		require.Equal(t, http.StatusUnauthorized, res.StatusCode)
	})
}

func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler, database.Workspace, database.TemplateVersion) {
	t.Helper()
	org := dbgen.Organization(t, db, database.Organization{})
	user := dbgen.User(t, db, database.User{
		Status: database.UserStatusActive,
	})
	_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
		UserID:         user.ID,
		OrganizationID: org.ID,
	})
	templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
		OrganizationID: org.ID,
		CreatedBy:      user.ID,
	})
	template := dbgen.Template(t, db, database.Template{
		OrganizationID:  org.ID,
		ActiveVersionID: templateVersion.ID,
		CreatedBy:       user.ID,
	})
	workspace := dbgen.Workspace(t, db, database.Workspace{
		OwnerID:        user.ID,
		OrganizationID: org.ID,
		TemplateID:     template.ID,
	})
	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
		OrganizationID: org.ID,
	})
	resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
		JobID: job.ID,
	})
	_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
		WorkspaceID:       workspace.ID,
		JobID:             job.ID,
		TemplateVersionID: templateVersion.ID,
	})
	_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
		ResourceID: resource.ID,
		AuthToken:  authToken,
	})

	req := httptest.NewRequest("GET", "/", nil)
	rtr := chi.NewRouter()
	rtr.Use(mw)
	rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
		_ = httpmw.WorkspaceAgent(r)
		rw.WriteHeader(http.StatusOK)
	})

	return req, rtr, workspace, templateVersion
}