mirror of https://github.com/databricks/cli.git
Attempt to reduce test flakiness on Windows
This commit is contained in:
parent
ca45e53f42
commit
d3e9701f27
|
@ -6,6 +6,7 @@ import (
|
|||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type execOption func(context.Context, *exec.Cmd) error
|
||||
|
@ -69,10 +70,40 @@ func WithStdoutWriter(dst io.Writer) execOption {
|
|||
}
|
||||
}
|
||||
|
||||
// safeMultiWriter is a thread-safe io.Writer that writes to multiple writers.
|
||||
// It is functionality equivalent to io.MultiWriter, but is safe for concurrent use.
|
||||
type safeMultiWriter struct {
|
||||
writers []io.Writer
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// newSafeMultiWriter creates a new safeMultiWriter that writes to the provided writers.
|
||||
func newSafeMultiWriter(writers ...io.Writer) *safeMultiWriter {
|
||||
return &safeMultiWriter{writers: writers}
|
||||
}
|
||||
|
||||
// Write implements the io.Writer interface for safeMultiWriter.
|
||||
func (t *safeMultiWriter) Write(p []byte) (n int, err error) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
||||
for _, w := range t.writers {
|
||||
n, err = w.Write(p)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if n != len(p) {
|
||||
err = io.ErrShortWrite
|
||||
return
|
||||
}
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func WithCombinedOutput(buf *bytes.Buffer) execOption {
|
||||
return func(_ context.Context, c *exec.Cmd) error {
|
||||
c.Stdout = io.MultiWriter(buf, c.Stdout)
|
||||
c.Stderr = io.MultiWriter(buf, c.Stderr)
|
||||
c.Stdout = newSafeMultiWriter(buf, c.Stdout)
|
||||
c.Stderr = newSafeMultiWriter(buf, c.Stderr)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue