Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,8 @@ These are methods on a pipe that change its configuration:
| [`WithReader`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithReader) | pipe source |
| [`WithStderr`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithStderr) | standard error output stream for command |
| [`WithStdout`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithStdout) | standard output stream for pipe |
| [`WithContext`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithContext) | context used in
exec and http requests |

## Filters

Expand Down
24 changes: 20 additions & 4 deletions script.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package script
import (
"bufio"
"container/ring"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
Expand Down Expand Up @@ -38,6 +39,7 @@ type Pipe struct {
err error
stderr io.Writer
env []string
ctx context.Context
}

// Args creates a pipe containing the program's command-line arguments from
Expand Down Expand Up @@ -174,6 +176,7 @@ func NewPipe() *Pipe {
stdout: os.Stdout,
httpClient: http.DefaultClient,
env: nil,
ctx: context.Background(),
}
}

Expand Down Expand Up @@ -432,7 +435,7 @@ func (p *Pipe) Exec(cmdLine string) *Pipe {
if err != nil {
return err
}
cmd := exec.Command(args[0], args[1:]...)
cmd := exec.CommandContext(p.ctx, args[0], args[1:]...)
cmd.Stdin = r
cmd.Stdout = w
cmd.Stderr = w
Expand Down Expand Up @@ -470,6 +473,9 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe {
return p.Filter(func(r io.Reader, w io.Writer) error {
scanner := newScanner(r)
for scanner.Scan() {
if p.ctx.Err() != nil {
return p.ctx.Err()
}
cmdLine := new(strings.Builder)
err := tpl.Execute(cmdLine, scanner.Text())
if err != nil {
Expand All @@ -479,7 +485,7 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe {
if err != nil {
return err
}
cmd := exec.Command(args[0], args[1:]...)
cmd := exec.CommandContext(p.ctx, args[0], args[1:]...)
cmd.Stdout = w
cmd.Stderr = w
pipeStderr := p.stdErr()
Expand Down Expand Up @@ -652,7 +658,7 @@ func (p *Pipe) Freq() *Pipe {
// the request body, and produces the server's response. See [Pipe.Do] for how
// the HTTP response status is interpreted.
func (p *Pipe) Get(url string) *Pipe {
req, err := http.NewRequest(http.MethodGet, url, p.Reader)
req, err := http.NewRequestWithContext(p.ctx, http.MethodGet, url, p.Reader)
if err != nil {
return p.WithError(err)
}
Expand Down Expand Up @@ -807,7 +813,7 @@ func (p *Pipe) MatchRegexp(re *regexp.Regexp) *Pipe {
// the request body, and produces the server's response. See [Pipe.Do] for how
// the HTTP response status is interpreted.
func (p *Pipe) Post(url string) *Pipe {
req, err := http.NewRequest(http.MethodPost, url, p.Reader)
req, err := http.NewRequestWithContext(p.ctx, http.MethodPost, url, p.Reader)
if err != nil {
return p.WithError(err)
}
Expand Down Expand Up @@ -1014,6 +1020,16 @@ func (p *Pipe) WithStdout(w io.Writer) *Pipe {
return p
}

// WithContext sets the context for subsequent [Pipe.Exec], [Pipe.ExecForEach], [Pipe.Get] and [Pipe.Post] commands .
func (p *Pipe) WithContext(ctx context.Context) *Pipe {
p.ctx = ctx
return p
}

func WithContext(ctx context.Context) *Pipe {
return NewPipe().WithContext(ctx)
}

// WriteFile writes the pipe's contents to the file path, truncating it if it
// exists, and returns the number of bytes successfully written, or an error.
func (p *Pipe) WriteFile(path string) (int64, error) {
Expand Down
102 changes: 102 additions & 0 deletions script_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package script_test
import (
"bufio"
"bytes"
"context"
"crypto/sha256"
"crypto/sha512"
"errors"
Expand All @@ -18,6 +19,7 @@ import (
"strings"
"testing"
"testing/iotest"
"time"

"github.com/bitfield/script"
"github.com/google/go-cmp/cmp"
Expand Down Expand Up @@ -2175,6 +2177,106 @@ func TestHashSums_OutputsEmptyStringForFileThatCannotBeHashed(t *testing.T) {
}
}

func TestWithContext_CallTimeoutMidstream(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
defer cancel()
err := script.WithContext(ctx).Exec("sleep 10").Wait()
if err == nil {
t.Fatal("expected error due to context cancellation, got nil")
}
}

func TestWithContext_CallCancelBeforeExec(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
cancel()
err := script.WithContext(ctx).Exec("echo Hello, World").Wait()
if err == nil {
t.Fatal("expected error due to context cancellation, got nil")
}
}

func TestWithContext_GetWithContextTimeout(t *testing.T) {
t.Parallel()
tcs := []struct {
name string
timeout time.Duration
expErr bool
}{
{
name: "timeout less than server response time",
timeout: 500 * time.Millisecond,
expErr: true,
},
{
name: "timeout greater than server response time",
timeout: 2 * time.Second,
expErr: false,
},
}
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
time.Sleep(1 * time.Second)
}))
defer ts.Close()

for _, tc := range tcs {
t.Run(tc.name, func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), tc.timeout)
defer cancel()
_, err := script.WithContext(ctx).Echo("request data").Get(ts.URL).String()
if tc.expErr {
if err == nil {
t.Fatalf("expected error due to context timeout, got nil")
}
} else {
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
}
})
}
}

func TestWithContext_PostWithContextTimeout(t *testing.T) {
t.Parallel()
tcs := []struct {
name string
timeout time.Duration
expErr bool
}{
{
name: "timeout less than server response time",
timeout: 500 * time.Millisecond,
expErr: true,
},
{
name: "timeout greater than server response time",
timeout: 2 * time.Second,
expErr: false,
},
}
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
time.Sleep(1 * time.Second)
}))
defer ts.Close()

for _, tc := range tcs {
t.Run(tc.name, func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), tc.timeout)
defer cancel()
_, err := script.WithContext(ctx).Echo("request data").Post(ts.URL).String()
if tc.expErr {
if err == nil {
t.Fatalf("expected error due to context timeout, got nil")
}
} else {
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
}
})
}
}

func ExampleArgs() {
script.Args().Stdout()
// prints command-line arguments
Expand Down