coder · Apr 20, 2023
diff --git a/‎cli/clibase/option.go
Lines changed: 11 additions & 0 deletions b/‎cli/clibase/option.go
Lines changed: 11 additions & 0 deletions
diff --git a/‎cli/clitest/clitest.go
Lines changed: 14 additions & 13 deletions b/‎cli/clitest/clitest.go
Lines changed: 14 additions & 13 deletions
@@ -80,6 +80,17 @@ func (s *OptionSet) Add(opts ...Option) {
 	*s = append(*s, opts...)
 }
 
+// Filter will only return options that match the given filter. (return true)
+func (s OptionSet) Filter(filter func(opt Option) bool) OptionSet {
+	cpy := make(OptionSet, 0)
+	for _, opt := range s {
+		if filter(opt) {
+			cpy = append(cpy, opt)
+		}
+	}
+	return cpy
+}
+
 // FlagSet returns a pflag.FlagSet for the OptionSet.
 func (s *OptionSet) FlagSet() *pflag.FlagSet {
 	if s == nil {
 
@@ -127,8 +127,8 @@ func extractTar(t *testing.T, data []byte, directory string) {
 	}
 }
 
-// Start runs the command in a goroutine and cleans it up when
-// the test completed.
+// Start runs the command in a goroutine and cleans it up when the test
+// completed.
 func Start(t *testing.T, inv *clibase.Invocation) {
 	t.Helper()
 
@@ -170,7 +170,7 @@ func (w *ErrorWaiter) Wait() error {
 		var ok bool
 		w.cachedError, ok = <-w.c
 		if !ok {
-			panic("unexpoected channel close")
+			panic("unexpected channel close")
 		}
 	})
 	return w.cachedError
@@ -196,18 +196,18 @@ func (w *ErrorWaiter) RequireAs(want interface{}) {
 	require.ErrorAs(w.t, w.Wait(), want)
 }
 
-// StartWithWaiter runs the command in a goroutine but returns the error
-// instead of asserting it. This is useful for testing error cases.
+// StartWithWaiter runs the command in a goroutine but returns the error instead
+// of asserting it. This is useful for testing error cases.
 func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
 	t.Helper()
 
-	errCh := make(chan error, 1)
-
-	var cleaningUp atomic.Bool
-
 	var (
 		ctx    = inv.Context()
 		cancel func()
+
+		cleaningUp atomic.Bool
+		errCh      = make(chan error, 1)
+		doneCh     = make(chan struct{})
 	)
 	if _, ok := ctx.Deadline(); !ok {
 		ctx, cancel = context.WithDeadline(ctx, time.Now().Add(testutil.WaitMedium))
@@ -218,12 +218,13 @@ func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
 	inv = inv.WithContext(ctx)
 
 	go func() {
+		defer close(doneCh)
 		defer close(errCh)
 		err := inv.Run()
 		if cleaningUp.Load() && errors.Is(err, context.DeadlineExceeded) {
-			// If we're cleaning up, this error is likely related to the
-			// CLI teardown process. E.g., the server could be slow to shut
-			// down Postgres.
+			// If we're cleaning up, this error is likely related to the CLI
+			// teardown process. E.g., the server could be slow to shut down
+			// Postgres.
 			t.Logf("command %q timed out during test cleanup", inv.Command.FullName())
 		}
 		// Whether or not this fails the test is left to the caller.
@@ -235,7 +236,7 @@ func StartWithWaiter(t *testing.T, inv *clibase.Invocation) *ErrorWaiter {
 	t.Cleanup(func() {
 		cancel()
 		cleaningUp.Store(true)
-		<-errCh
+		<-doneCh
 	})
 	return &ErrorWaiter{c: errCh, t: t}
 }