From 36d75c241d163227ec86567f171903342b5d7287 Mon Sep 17 00:00:00 2001 From: Vasilis Vamvakas Date: Mon, 6 Jul 2026 15:40:45 +0100 Subject: [PATCH] pass optional context to pipe --- README.md | 2 + script.go | 24 ++++++++++-- script_test.go | 102 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 038e31e..c80efc6 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/script.go b/script.go index 2195da4..fd1c045 100644 --- a/script.go +++ b/script.go @@ -3,6 +3,7 @@ package script import ( "bufio" "container/ring" + "context" "crypto/sha256" "encoding/base64" "encoding/hex" @@ -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 @@ -174,6 +176,7 @@ func NewPipe() *Pipe { stdout: os.Stdout, httpClient: http.DefaultClient, env: nil, + ctx: context.Background(), } } @@ -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 @@ -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 { @@ -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() @@ -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) } @@ -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) } @@ -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) { diff --git a/script_test.go b/script_test.go index 66e4df5..383a12b 100644 --- a/script_test.go +++ b/script_test.go @@ -3,6 +3,7 @@ package script_test import ( "bufio" "bytes" + "context" "crypto/sha256" "crypto/sha512" "errors" @@ -18,6 +19,7 @@ import ( "strings" "testing" "testing/iotest" + "time" "github.com/bitfield/script" "github.com/google/go-cmp/cmp" @@ -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