storj/internal/testcontext/context.go
2018-10-31 20:55:14 +02:00

105 lines
2.0 KiB
Go

// Copyright (C) 2018 Storj Labs, Inc.
// See LICENSE for copying information.
package testcontext
import (
"context"
"io/ioutil"
"os"
"path/filepath"
"sync"
"testing"
"golang.org/x/sync/errgroup"
)
// Context is a context that has utility methods for testing and waiting for asynchronous errors.
type Context struct {
context.Context
group *errgroup.Group
test testing.TB
once sync.Once
directory string
}
// New creates a new test context
func New(test testing.TB) *Context {
group, ctx := errgroup.WithContext(context.Background())
return &Context{
Context: ctx,
group: group,
test: test,
}
}
// Go runs fn in a goroutine.
// Call Wait to check the result
func (ctx *Context) Go(fn func() error) {
ctx.test.Helper()
ctx.group.Go(fn)
}
// 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
ctx.directory, err = ioutil.TempDir("", ctx.test.Name())
if err != nil {
ctx.test.Fatal(err)
}
})
dir := filepath.Join(append([]string{ctx.directory}, subs...)...)
_ = os.MkdirAll(dir, 0644)
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()
err := ctx.group.Wait()
if err != nil {
ctx.test.Fatal(err)
}
}
// deleteTemporary tries to delete temporary directory
func (ctx *Context) deleteTemporary() {
if ctx.directory == "" {
return
}
err := os.RemoveAll(ctx.directory)
if err != nil {
ctx.test.Fatal(err)
}
ctx.directory = ""
}