diff --git a/internal/testutil/retry.go b/internal/testutil/retry.go new file mode 100644 index 00000000000..7308d7364cf --- /dev/null +++ b/internal/testutil/retry.go @@ -0,0 +1,116 @@ +// Copyright 2019 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package testutil + +import ( + "bytes" + "fmt" + "path/filepath" + "runtime" + "strconv" + "testing" + "time" +) + +// Retry runs function f for up to maxAttempts times until f returns successfully, and reports whether f was run successfully. +// It will sleep for the given period between invocations of f. +// Use the provided *testutil.R instead of a *testing.T from the function. +func Retry(t *testing.T, maxAttempts int, sleep time.Duration, f func(r *R)) bool { + for attempt := 1; attempt <= maxAttempts; attempt++ { + r := &R{Attempt: attempt, log: &bytes.Buffer{}} + + f(r) + + if !r.failed { + if r.log.Len() != 0 { + t.Logf("Success after %d attempts:%s", attempt, r.log.String()) + } + return true + } + + if attempt == maxAttempts { + t.Logf("FAILED after %d attempts:%s", attempt, r.log.String()) + t.Fail() + } + + time.Sleep(sleep) + } + return false +} + +// RetryWithoutTest is a variant of Retry that does not use a testing parameter. +// It is meant for testing utilities that do not pass around the testing context, such as cloudrunci. +func RetryWithoutTest(maxAttempts int, sleep time.Duration, f func(r *R)) bool { + for attempt := 1; attempt <= maxAttempts; attempt++ { + r := &R{Attempt: attempt, log: &bytes.Buffer{}} + + f(r) + + if !r.failed { + if r.log.Len() != 0 { + r.Logf("Success after %d attempts:%s", attempt, r.log.String()) + } + return true + } + + if attempt == maxAttempts { + r.Logf("FAILED after %d attempts:%s", attempt, r.log.String()) + return false + } + + time.Sleep(sleep) + } + return false +} + +// R is passed to each run of a flaky test run, manages state and accumulates log statements. +type R struct { + // The number of current attempt. + Attempt int + + failed bool + log *bytes.Buffer +} + +// Fail marks the run as failed, and will retry once the function returns. +func (r *R) Fail() { + r.failed = true +} + +// Errorf is equivalent to Logf followed by Fail. +func (r *R) Errorf(s string, v ...interface{}) { + r.logf(s, v...) + r.Fail() +} + +// Logf formats its arguments and records it in the error log. +// The text is only printed for the final unsuccessful run or the first successful run. +func (r *R) Logf(s string, v ...interface{}) { + r.logf(s, v...) +} + +func (r *R) logf(s string, v ...interface{}) { + fmt.Fprint(r.log, "\n") + fmt.Fprint(r.log, lineNumber()) + fmt.Fprintf(r.log, s, v...) +} + +func lineNumber() string { + _, file, line, ok := runtime.Caller(3) // logf, public func, user function + if !ok { + return "" + } + return filepath.Base(file) + ":" + strconv.Itoa(line) + ": " +} diff --git a/internal/testutil/retry_test.go b/internal/testutil/retry_test.go new file mode 100644 index 00000000000..170e7abf4c0 --- /dev/null +++ b/internal/testutil/retry_test.go @@ -0,0 +1,76 @@ +// Copyright 2019 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package testutil + +import ( + "testing" + "time" +) + +func TestRetry(t *testing.T) { + Retry(t, 5, time.Millisecond, func(r *R) { + if r.Attempt == 2 { + return + } + r.Fail() + }) +} + +func TestRetryAttempts(t *testing.T) { + var attempts int + Retry(t, 10, time.Millisecond, func(r *R) { + r.Logf("This line should appear only once.") + r.Logf("attempt=%d", r.Attempt) + attempts = r.Attempt + + // Retry 5 times. + if r.Attempt == 5 { + return + } + r.Fail() + }) + + if attempts != 5 { + t.Errorf("attempts=%d; want %d", attempts, 5) + } +} + +func TestRetryWithoutTest(t *testing.T) { + RetryWithoutTest(5, time.Millisecond, func(r *R) { + if r.Attempt == 2 { + return + } + r.Fail() + }) +} + +func TestRetryWithoutTestAttempts(t *testing.T) { + var attempts int + RetryWithoutTest(10, time.Millisecond, func(r *R) { + r.Logf("This line should appear only once.") + r.Logf("attempt=%d", r.Attempt) + attempts = r.Attempt + + // Retry 5 times. + if r.Attempt == 5 { + return + } + r.Fail() + }) + + if attempts != 5 { + t.Errorf("attempts=%d; want %d", attempts, 5) + } +}