coder · ammario · Mar 23, 2023 · Mar 6, 2023 · Mar 6, 2023 · Mar 6, 2023
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -302,7 +302,7 @@ jobs:
             echo "cover=false" >> $GITHUB_OUTPUT
           fi
 
-          gotestsum --junitfile="gotests.xml" --packages="./..." -- -parallel=8 -timeout=5m -short -failfast $COVERAGE_FLAGS
+          gotestsum --junitfile="gotests.xml" --packages="./..." -- -parallel=8 -timeout=7m -short -failfast $COVERAGE_FLAGS
 
       - uses: actions/upload-artifact@v3
         if: success() || failure()

diff --git a/Makefile b/Makefile
@@ -501,8 +501,6 @@ docs/admin/prometheus.md: scripts/metricsdocgen/main.go scripts/metricsdocgen/me
 	yarn run format:write:only ../docs/admin/prometheus.md
 
 docs/cli.md: scripts/clidocgen/main.go $(GO_SRC_FILES) docs/manifest.json
-	# TODO(@ammario): re-enable server.md once we finish clibase migration.
-	ls ./docs/cli/*.md | grep -vP "\/coder_server" | xargs rm
 	BASE_PATH="." go run ./scripts/clidocgen
 	cd site
 	yarn run format:write:only ../docs/cli.md ../docs/cli/*.md ../docs/manifest.json
@@ -519,7 +517,7 @@ coderd/apidoc/swagger.json: $(shell find ./scripts/apidocgen $(FIND_EXCLUSIONS)
 update-golden-files: cli/testdata/.gen-golden helm/tests/testdata/.gen-golden
 .PHONY: update-golden-files
 
-cli/testdata/.gen-golden: $(wildcard cli/testdata/*.golden) $(GO_SRC_FILES)
+cli/testdata/.gen-golden: $(wildcard cli/testdata/*.golden) $(wildcard cli/*.tpl) $(GO_SRC_FILES)
 	go test ./cli -run=TestCommandHelp -update
 	touch "$@"
 

diff --git a/cli/agent.go b/cli/agent.go
@@ -16,7 +16,6 @@ import (
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
-	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 	"gopkg.in/natefinch/lumberjack.v2"
 
@@ -25,34 +24,27 @@ import (
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/agent/reaper"
 	"github.com/coder/coder/buildinfo"
-	"github.com/coder/coder/cli/cliflag"
+	"github.com/coder/coder/cli/clibase"
 	"github.com/coder/coder/codersdk/agentsdk"
 )
 
-func workspaceAgent() *cobra.Command {
+func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 	var (
 		auth          string
 		logDir        string
 		pprofAddress  string
 		noReap        bool
 		sshMaxTimeout time.Duration
 	)
-	cmd := &cobra.Command{
-		Use: "agent",
+	cmd := &clibase.Cmd{
+		Use:   "agent",
+		Short: `Starts the Coder workspace agent.`,
 		// This command isn't useful to manually execute.
 		Hidden: true,
-		RunE: func(cmd *cobra.Command, _ []string) error {
-			ctx, cancel := context.WithCancel(cmd.Context())
+		Handler: func(inv *clibase.Invocation) error {
+			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
 
-			rawURL, err := cmd.Flags().GetString(varAgentURL)
-			if err != nil {
-				return xerrors.Errorf("CODER_AGENT_URL must be set: %w", err)
-			}
-			coderURL, err := url.Parse(rawURL)
-			if err != nil {
-				return xerrors.Errorf("parse %q: %w", rawURL, err)
-			}
 			agentPorts := map[int]string{}
 
 			isLinux := runtime.GOOS == "linux"
@@ -65,7 +57,7 @@ func workspaceAgent() *cobra.Command {
 					MaxSize:  5, // MB
 				}
 				defer logWriter.Close()
-				logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
+				logger := slog.Make(sloghuman.Sink(inv.Stderr), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
 
 				logger.Info(ctx, "spawning reaper process")
 				// Do not start a reaper on the child process. It's important
@@ -107,15 +99,15 @@ func workspaceAgent() *cobra.Command {
 			logWriter := &closeWriter{w: ljLogger}
 			defer logWriter.Close()
 
-			logger := slog.Make(sloghuman.Sink(cmd.ErrOrStderr()), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
+			logger := slog.Make(sloghuman.Sink(inv.Stderr), sloghuman.Sink(logWriter)).Leveled(slog.LevelDebug)
 
 			version := buildinfo.Version()
 			logger.Info(ctx, "starting agent",
-				slog.F("url", coderURL),
+				slog.F("url", r.agentURL),
 				slog.F("auth", auth),
 				slog.F("version", version),
 			)
-			client := agentsdk.New(coderURL)
+			client := agentsdk.New(r.agentURL)
 			client.SDK.Logger = logger
 			// Set a reasonable timeout so requests can't hang forever!
 			// The timeout needs to be reasonably long, because requests
@@ -139,7 +131,7 @@ func workspaceAgent() *cobra.Command {
 			var exchangeToken func(context.Context) (agentsdk.AuthenticateResponse, error)
 			switch auth {
 			case "token":
-				token, err := cmd.Flags().GetString(varAgentToken)
+				token, err := inv.ParsedFlags().GetString(varAgentToken)
 				if err != nil {
 					return xerrors.Errorf("CODER_AGENT_TOKEN must be set for token auth: %w", err)
 				}
@@ -220,11 +212,44 @@ func workspaceAgent() *cobra.Command {
 		},
 	}
 
-	cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AGENT_AUTH", "token", "Specify the authentication type to use for the agent")
-	cliflag.StringVarP(cmd.Flags(), &logDir, "log-dir", "", "CODER_AGENT_LOG_DIR", os.TempDir(), "Specify the location for the agent log files")
-	cliflag.StringVarP(cmd.Flags(), &pprofAddress, "pprof-address", "", "CODER_AGENT_PPROF_ADDRESS", "127.0.0.1:6060", "The address to serve pprof.")
-	cliflag.BoolVarP(cmd.Flags(), &noReap, "no-reap", "", "", false, "Do not start a process reaper.")
-	cliflag.DurationVarP(cmd.Flags(), &sshMaxTimeout, "ssh-max-timeout", "", "CODER_AGENT_SSH_MAX_TIMEOUT", time.Duration(0), "Specify the max timeout for a SSH connection")
+	cmd.Options = clibase.OptionSet{
+		{
+			Flag:        "auth",
+			Default:     "token",
+			Description: "Specify the authentication type to use for the agent.",
+			Env:         "CODER_AGENT_AUTH",
+			Value:       clibase.StringOf(&auth),
+		},
+		{
+			Flag:        "log-dir",
+			Default:     os.TempDir(),
+			Description: "Specify the location for the agent log files.",
+			Env:         "CODER_AGENT_LOG_DIR",
+			Value:       clibase.StringOf(&logDir),
+		},
+		{
+			Flag:        "pprof-address",
+			Default:     "127.0.0.1:6060",
+			Env:         "CODER_AGENT_PPROF_ADDRESS",
+			Value:       clibase.StringOf(&pprofAddress),
+			Description: "The address to serve pprof.",
+		},
+		{
+			Flag: "no-reap",
+
+			Env:         "",
+			Description: "Do not start a process reaper.",
+			Value:       clibase.BoolOf(&noReap),
+		},
+		{
+			Flag:        "ssh-max-timeout",
+			Default:     "0",
+			Env:         "CODER_AGENT_SSH_MAX_TIMEOUT",
+			Description: "Specify the max timeout for a SSH connection.",
+			Value:       clibase.DurationOf(&sshMaxTimeout),
+		},
+	}
+
 	return cmd
 }
 

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -16,7 +16,7 @@ import (
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
-	"github.com/coder/coder/testutil"
+	"github.com/coder/coder/pty/ptytest"
 )
 
 func TestWorkspaceAgent(t *testing.T) {
@@ -40,24 +40,20 @@ func TestWorkspaceAgent(t *testing.T) {
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
 		logDir := t.TempDir()
-		cmd, _ := clitest.New(t,
+		inv, _ := clitest.New(t,
 			"agent",
 			"--auth", "token",
 			"--agent-token", authToken,
 			"--agent-url", client.URL.String(),
 			"--log-dir", logDir,
 		)
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
-		defer cancel()
-		errC := make(chan error, 1)
-		go func() {
-			errC <- cmd.ExecuteContext(ctx)
-		}()
-		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 
-		cancel()
-		err := <-errC
-		require.NoError(t, err)
+		pty := ptytest.New(t).Attach(inv)
+
+		clitest.Start(t, inv)
+		pty.ExpectMatch("starting agent")
+
+		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 
 		info, err := os.Stat(filepath.Join(logDir, "coder-agent.log"))
 		require.NoError(t, err)
@@ -96,16 +92,14 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
-		cmd, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String())
+		inv, _ := clitest.New(t, "agent", "--auth", "azure-instance-identity", "--agent-url", client.URL.String())
+		inv = inv.WithContext(
+			//nolint:revive,staticcheck
+			context.WithValue(inv.Context(), "azure-client", metadataClient),
+		)
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
-		errC := make(chan error)
-		go func() {
-			// A linting error occurs for weakly typing the context value here.
-			//nolint // The above seems reasonable for a one-off test.
-			ctx := context.WithValue(ctx, "azure-client", metadataClient)
-			errC <- cmd.ExecuteContext(ctx)
-		}()
+		clitest.Start(t, inv)
 		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 		workspace, err := client.Workspace(ctx, workspace.ID)
 		require.NoError(t, err)
@@ -117,9 +111,6 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.NoError(t, err)
 		defer dialer.Close()
 		require.True(t, dialer.AwaitReachable(context.Background()))
-		cancelFunc()
-		err = <-errC
-		require.NoError(t, err)
 	})
 
 	t.Run("AWS", func(t *testing.T) {
@@ -154,36 +145,29 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
-		cmd, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String())
-		ctx, cancelFunc := context.WithCancel(context.Background())
-		defer cancelFunc()
-		errC := make(chan error)
-		go func() {
-			// A linting error occurs for weakly typing the context value here.
-			//nolint // The above seems reasonable for a one-off test.
-			ctx := context.WithValue(ctx, "aws-client", metadataClient)
-			errC <- cmd.ExecuteContext(ctx)
-		}()
+		inv, _ := clitest.New(t, "agent", "--auth", "aws-instance-identity", "--agent-url", client.URL.String())
+		inv = inv.WithContext(
+			//nolint:revive,staticcheck
+			context.WithValue(inv.Context(), "aws-client", metadataClient),
+		)
+		clitest.Start(t, inv)
 		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
-		workspace, err := client.Workspace(ctx, workspace.ID)
+		workspace, err := client.Workspace(inv.Context(), workspace.ID)
 		require.NoError(t, err)
 		resources := workspace.LatestBuild.Resources
 		if assert.NotEmpty(t, resources) && assert.NotEmpty(t, resources[0].Agents) {
 			assert.NotEmpty(t, resources[0].Agents[0].Version)
 		}
-		dialer, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, nil)
+		dialer, err := client.DialWorkspaceAgent(inv.Context(), resources[0].Agents[0].ID, nil)
 		require.NoError(t, err)
 		defer dialer.Close()
 		require.True(t, dialer.AwaitReachable(context.Background()))
-		cancelFunc()
-		err = <-errC
-		require.NoError(t, err)
 	})
 
 	t.Run("GoogleCloud", func(t *testing.T) {
 		t.Parallel()
 		instanceID := "instanceidentifier"
-		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
+		validator, metadataClient := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator:     validator,
 			IncludeProvisionerDaemon: true,
@@ -212,16 +196,18 @@ func TestWorkspaceAgent(t *testing.T) {
 		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
 		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 
-		cmd, _ := clitest.New(t, "agent", "--auth", "google-instance-identity", "--agent-url", client.URL.String())
-		ctx, cancelFunc := context.WithCancel(context.Background())
-		defer cancelFunc()
-		errC := make(chan error)
-		go func() {
-			// A linting error occurs for weakly typing the context value here.
-			//nolint // The above seems reasonable for a one-off test.
-			ctx := context.WithValue(ctx, "gcp-client", metadata)
-			errC <- cmd.ExecuteContext(ctx)
-		}()
+		inv, cfg := clitest.New(t, "agent", "--auth", "google-instance-identity", "--agent-url", client.URL.String())
+		ptytest.New(t).Attach(inv)
+		clitest.SetupConfig(t, client, cfg)
+		clitest.Start(t,
+			inv.WithContext(
+				//nolint:revive,staticcheck
+				context.WithValue(context.Background(), "gcp-client", metadataClient),
+			),
+		)
+
+		ctx := inv.Context()
+
 		coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
 		workspace, err := client.Workspace(ctx, workspace.ID)
 		require.NoError(t, err)
@@ -248,9 +234,5 @@ func TestWorkspaceAgent(t *testing.T) {
 		require.NoError(t, err)
 		_, err = uuid.Parse(strings.TrimSpace(string(token)))
 		require.NoError(t, err)
-
-		cancelFunc()
-		err = <-errC
-		require.NoError(t, err)
 	})
 }
diff --git a/cli/clibase/clibase.go b/cli/clibase/clibase.go
@@ -1,10 +1,6 @@
 // Package clibase offers an all-in-one solution for a highly configurable CLI
-// application. Within Coder, we use it for our `server` subcommand, which
-// demands more functionality than cobra/viper can offer.
-//
-// We will extend its usage to the rest of our application, completely replacing
-// cobra/viper. It's also a candidate to be broken out into its own open-source
-// library, so we avoid deep coupling with Coder concepts.
+// application. Within Coder, we use it for all of our subcommands, which
+// demands more functionality than cobra/viber offers.
 //
 // The Command interface is loosely based on the chi middleware pattern and
 // http.Handler/HandlerFunc.