storj/internal/testcontext/context.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 = ""
}