Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
Comment thread
billvamva marked this conversation as resolved.
Outdated
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 .
Comment thread
billvamva marked this conversation as resolved.
Outdated
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)
Comment thread
billvamva marked this conversation as resolved.
Outdated
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)
Comment thread
billvamva marked this conversation as resolved.
Outdated
cancel()
err := script.WithContext(ctx).Exec("echo Hello, World").Wait()
Comment thread
billvamva marked this conversation as resolved.
Outdated
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)
Comment thread
billvamva marked this conversation as resolved.
Outdated
}))
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",
Comment thread
billvamva marked this conversation as resolved.
Outdated
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