storj/internal/testcontext/context.go

206 lines
4.1 KiB
Go
Raw Normal View History

2019-01-24 20:15:10 +00:00
// Copyright (C) 2019 Storj Labs, Inc.
2018-10-29 14:16:36 +00:00
// See LICENSE for copying information.
package testcontext
import (
"context"
"fmt"
2018-10-29 14:16:36 +00:00
"io/ioutil"
"os"
"path/filepath"
"regexp"
"runtime"
"strings"
2018-10-29 14:16:36 +00:00
"sync"
"time"
2018-10-29 14:16:36 +00:00
"golang.org/x/sync/errgroup"
"storj.io/storj/internal/memory"
2018-10-29 14:16:36 +00:00
)
const defaultTimeout = 3 * time.Minute
2018-10-29 14:16:36 +00:00
// Context is a context that has utility methods for testing and waiting for asynchronous errors.
type Context struct {
context.Context
timedctx context.Context
cancel context.CancelFunc
2018-10-29 14:16:36 +00:00
group *errgroup.Group
test TB
2018-10-29 14:16:36 +00:00
once sync.Once
directory string
mu sync.Mutex
running []caller
}
type caller struct {
pc uintptr
file string
line int
ok bool
done bool
}
// TB is a subset of testing.TB methods
type TB interface {
Name() string
Helper()
Error(args ...interface{})
Fatal(args ...interface{})
2018-10-29 14:16:36 +00:00
}
// New creates a new test context
func New(test TB) *Context {
return NewWithTimeout(test, defaultTimeout)
}
// NewWithTimeout creates a new test context with a given timeout
func NewWithTimeout(test TB, timeout time.Duration) *Context {
timedctx, cancel := context.WithTimeout(context.Background(), timeout)
group, errctx := errgroup.WithContext(timedctx)
ctx := &Context{
Context: errctx,
timedctx: timedctx,
cancel: cancel,
group: group,
test: test,
2018-10-29 14:16:36 +00:00
}
return ctx
2018-10-29 14:16:36 +00:00
}
// Go runs fn in a goroutine.
// Call Wait to check the result
func (ctx *Context) Go(fn func() error) {
ctx.test.Helper()
pc, file, line, ok := runtime.Caller(1)
ctx.mu.Lock()
index := len(ctx.running)
ctx.running = append(ctx.running, caller{pc, file, line, ok, false})
ctx.mu.Unlock()
ctx.group.Go(func() error {
defer func() {
ctx.mu.Lock()
ctx.running[index].done = true
ctx.mu.Unlock()
}()
return fn()
})
2018-10-29 14:16:36 +00:00
}
// Check calls fn and checks result
func (ctx *Context) Check(fn func() error) {
ctx.test.Helper()
err := fn()
if err != nil {
ctx.test.Fatal(err)
}
}
// Dir returns a directory path inside temp
func (ctx *Context) Dir(subs ...string) string {
ctx.test.Helper()
ctx.once.Do(func() {
var err error
pattern := regexp.MustCompile(`[\\/]`)
ctx.directory, err = ioutil.TempDir("", pattern.ReplaceAllString(ctx.test.Name(), "_"))
2018-10-29 14:16:36 +00:00
if err != nil {
ctx.test.Fatal(err)
}
})
dir := filepath.Join(append([]string{ctx.directory}, subs...)...)
_ = os.MkdirAll(dir, 0744)
2018-10-29 14:16:36 +00:00
return dir
}
// File returns a filepath inside temp
func (ctx *Context) File(subs ...string) string {
ctx.test.Helper()
if len(subs) == 0 {
ctx.test.Fatal("expected more than one argument")
}
dir := ctx.Dir(subs[:len(subs)-1]...)
return filepath.Join(dir, subs[len(subs)-1])
}
// Cleanup waits everything to be completed,
// checks errors and tries to cleanup directories
func (ctx *Context) Cleanup() {
ctx.test.Helper()
defer ctx.deleteTemporary()
defer ctx.cancel()
alldone := make(chan error, 1)
go func() {
alldone <- ctx.group.Wait()
defer close(alldone)
}()
select {
case <-ctx.timedctx.Done():
ctx.reportRunning()
case err := <-alldone:
if err != nil {
ctx.test.Fatal(err)
}
}
}
func (ctx *Context) reportRunning() {
ctx.mu.Lock()
defer ctx.mu.Unlock()
var problematic []caller
for _, caller := range ctx.running {
if !caller.done {
problematic = append(problematic, caller)
}
}
var message strings.Builder
message.WriteString("Test exceeded timeout")
if len(problematic) > 0 {
message.WriteString("\nsome goroutines are still running, did you forget to shut them down?")
for _, caller := range problematic {
fnname := ""
if fn := runtime.FuncForPC(caller.pc); fn != nil {
fnname = fn.Name()
}
fmt.Fprintf(&message, "\n%s:%d: %s", caller.file, caller.line, fnname)
}
2018-10-29 14:16:36 +00:00
}
ctx.test.Error(message.String())
stack := make([]byte, 1*memory.MB.Int())
n := runtime.Stack(stack, true)
stack = stack[:n]
ctx.test.Error("Full Stack Trace:\n", string(stack))
2018-10-29 14:16:36 +00:00
}
// deleteTemporary tries to delete temporary directory
func (ctx *Context) deleteTemporary() {
if ctx.directory == "" {
return
}
2018-10-29 14:16:36 +00:00
err := os.RemoveAll(ctx.directory)
if err != nil {
ctx.test.Fatal(err)
}
ctx.directory = ""
2018-10-29 14:16:36 +00:00
}