100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
|
package wait
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"io"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
tcexec "github.com/testcontainers/testcontainers-go/exec"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Implement interface
|
||
|
|
var _ Strategy = (*ExecStrategy)(nil)
|
||
|
|
var _ StrategyTimeout = (*ExecStrategy)(nil)
|
||
|
|
|
||
|
|
type ExecStrategy struct {
|
||
|
|
// all Strategies should have a startupTimeout to avoid waiting infinitely
|
||
|
|
timeout *time.Duration
|
||
|
|
cmd []string
|
||
|
|
|
||
|
|
// additional properties
|
||
|
|
ExitCodeMatcher func(exitCode int) bool
|
||
|
|
ResponseMatcher func(body io.Reader) bool
|
||
|
|
PollInterval time.Duration
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewExecStrategy constructs an Exec strategy ...
|
||
|
|
func NewExecStrategy(cmd []string) *ExecStrategy {
|
||
|
|
return &ExecStrategy{
|
||
|
|
cmd: cmd,
|
||
|
|
ExitCodeMatcher: defaultExitCodeMatcher,
|
||
|
|
ResponseMatcher: func(body io.Reader) bool { return true },
|
||
|
|
PollInterval: defaultPollInterval(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func defaultExitCodeMatcher(exitCode int) bool {
|
||
|
|
return exitCode == 0
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithStartupTimeout can be used to change the default startup timeout
|
||
|
|
func (ws *ExecStrategy) WithStartupTimeout(startupTimeout time.Duration) *ExecStrategy {
|
||
|
|
ws.timeout = &startupTimeout
|
||
|
|
return ws
|
||
|
|
}
|
||
|
|
|
||
|
|
func (ws *ExecStrategy) WithExitCodeMatcher(exitCodeMatcher func(exitCode int) bool) *ExecStrategy {
|
||
|
|
ws.ExitCodeMatcher = exitCodeMatcher
|
||
|
|
return ws
|
||
|
|
}
|
||
|
|
|
||
|
|
func (ws *ExecStrategy) WithResponseMatcher(matcher func(body io.Reader) bool) *ExecStrategy {
|
||
|
|
ws.ResponseMatcher = matcher
|
||
|
|
return ws
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithPollInterval can be used to override the default polling interval of 100 milliseconds
|
||
|
|
func (ws *ExecStrategy) WithPollInterval(pollInterval time.Duration) *ExecStrategy {
|
||
|
|
ws.PollInterval = pollInterval
|
||
|
|
return ws
|
||
|
|
}
|
||
|
|
|
||
|
|
// ForExec is a convenience method to assign ExecStrategy
|
||
|
|
func ForExec(cmd []string) *ExecStrategy {
|
||
|
|
return NewExecStrategy(cmd)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (ws *ExecStrategy) Timeout() *time.Duration {
|
||
|
|
return ws.timeout
|
||
|
|
}
|
||
|
|
|
||
|
|
func (ws *ExecStrategy) WaitUntilReady(ctx context.Context, target StrategyTarget) error {
|
||
|
|
timeout := defaultStartupTimeout()
|
||
|
|
if ws.timeout != nil {
|
||
|
|
timeout = *ws.timeout
|
||
|
|
}
|
||
|
|
|
||
|
|
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
return ctx.Err()
|
||
|
|
case <-time.After(ws.PollInterval):
|
||
|
|
exitCode, resp, err := target.Exec(ctx, ws.cmd, tcexec.Multiplexed())
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if !ws.ExitCodeMatcher(exitCode) {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
if ws.ResponseMatcher != nil && !ws.ResponseMatcher(resp) {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|