Skip to content

Commit

Permalink
Make Cobra runner compatible with testing interactive flows (#957)
Browse files Browse the repository at this point in the history
## Changes
This PR enables testing commands with stdin

## Tests
#914
  • Loading branch information
nfx authored Nov 7, 2023
1 parent 283f241 commit 56bcb6f
Showing 1 changed file with 49 additions and 1 deletion.
50 changes: 49 additions & 1 deletion internal/helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"math/rand"
Expand All @@ -17,6 +18,7 @@ import (

"github.com/databricks/cli/cmd"
_ "github.com/databricks/cli/cmd/version"
"github.com/databricks/cli/libs/cmdio"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/service/compute"
Expand Down Expand Up @@ -63,6 +65,8 @@ type cobraTestRunner struct {
args []string
stdout bytes.Buffer
stderr bytes.Buffer
stdinR *io.PipeReader
stdinW *io.PipeWriter

ctx context.Context

Expand Down Expand Up @@ -119,15 +123,46 @@ func (t *cobraTestRunner) registerFlagCleanup(c *cobra.Command) {
})
}

// Like [cobraTestRunner.Eventually], but more specific
func (t *cobraTestRunner) WaitForTextPrinted(text string, timeout time.Duration) {
t.Eventually(func() bool {
currentStdout := t.stdout.String()
return strings.Contains(currentStdout, text)
}, timeout, 50*time.Millisecond)
}

func (t *cobraTestRunner) WithStdin() {
reader, writer := io.Pipe()
t.stdinR = reader
t.stdinW = writer
}

func (t *cobraTestRunner) CloseStdin() {
if t.stdinW == nil {
panic("no standard input configured")
}
t.stdinW.Close()
}

func (t *cobraTestRunner) SendText(text string) {
if t.stdinW == nil {
panic("no standard input configured")
}
t.stdinW.Write([]byte(text + "\n"))
}

func (t *cobraTestRunner) RunBackground() {
var stdoutR, stderrR io.Reader
var stdoutW, stderrW io.WriteCloser
stdoutR, stdoutW = io.Pipe()
stderrR, stderrW = io.Pipe()
root := cmd.New(context.Background())
root := cmd.New(t.ctx)
root.SetOut(stdoutW)
root.SetErr(stderrW)
root.SetArgs(t.args)
if t.stdinW != nil {
root.SetIn(t.stdinR)
}

// Register cleanup function to restore flags to their original values
// once test has been executed. This is needed because flag values reside
Expand Down Expand Up @@ -239,6 +274,19 @@ func (c *cobraTestRunner) Eventually(condition func() bool, waitFor time.Duratio
}
}

func (t *cobraTestRunner) RunAndExpectOutput(heredoc string) {
stdout, _, err := t.Run()
require.NoError(t, err)
require.Equal(t, cmdio.Heredoc(heredoc), strings.TrimSpace(stdout.String()))
}

func (t *cobraTestRunner) RunAndParseJSON(v any) {
stdout, _, err := t.Run()
require.NoError(t, err)
err = json.Unmarshal(stdout.Bytes(), &v)
require.NoError(t, err)
}

func NewCobraTestRunner(t *testing.T, args ...string) *cobraTestRunner {
return &cobraTestRunner{
T: t,
Expand Down

0 comments on commit 56bcb6f

Please sign in to comment.