trace.go

package htest

import (
	"bytes"
	"fmt"
	"runtime"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/fatih/color"
)

// StackFrames is a list of StackFrame
type StackFrames []*StackFrame

// Trace parses the stacktrace and returns a slice of StackFrame
func Trace() (frames StackFrames) {
	for i := 0; ; i++ {
		pc, file, _, ok := runtime.Caller(i)
		if !ok {
			return frames
		}

		if file == "<autogenerated>" {
			continue
		}

		frame := NewStackFrame(pc)
		if frame != nil {
			frames = append(frames, frame)
		}
	}

}

// String returns a pretty string for each StackFrame
func (frames StackFrames) String() string {
	var buf bytes.Buffer
	for _, frame := range frames {
		buf.WriteString(frame.String() + "\n\r\t\t\t")
	}
	return buf.String()
}

// OnlyTests filters only the stackframes that are Tests
func (frames StackFrames) OnlyTests() (out StackFrames) {
	for _, frame := range frames {
		if frame.IsTest() {
			out = append(out, frame)
		}
	}
	return out
}

// StackFrame describes a single stackframe with its Program Counter, Function name, Package, File and Line number
type StackFrame struct {
	PC       uintptr
	File     string
	FuncName string
	Package  string
	Line     int
}

// NewStackFrame creates a new StackFrame instance
// It extracts the File, Line, Package and Function name for the passed Program Counter
func NewStackFrame(pc uintptr) *StackFrame {
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		return nil
	}
	pkg, funcName := extractPackageAndLine(fn.Name())
	file, line := fn.FileLine(pc - 1)

	s := &StackFrame{
		PC:       pc,
		File:     file,
		Line:     line,
		Package:  pkg,
		FuncName: funcName,
	}

	return s
}

var pkgColor = color.New(color.Italic, color.FgHiBlack).SprintFunc()
var gray = color.New(color.FgHiBlack).SprintFunc()
var funcColor = color.New(color.Bold).SprintFunc()
var lineColor = color.New(color.Bold, color.FgHiBlue).SprintFunc()

// String return a strings with the following format:
// 		example.com/mypackage.TestMyFunction:99
func (s *StackFrame) String() string {
	return fmt.Sprintf("%s.%s:%s", pkgColor(s.Package), funcColor(s.FuncName), funcColor(s.Line))
}

// IsTest determines if the found function is a Test
func (s *StackFrame) IsTest() bool {
	return isTest(s.FuncName, "Test") || isTest(s.FuncName, "Benchmark") || isTest(s.FuncName, "Example")
}

func extractPackageAndLine(funcPC string) (pkg, fn string) {
	fn = funcPC
	lastSlash := strings.LastIndex(funcPC, "/")
	if lastSlash >= 0 {
		pkg += funcPC[:lastSlash] + "/"
		fn = fn[lastSlash+1:]
	}
	if dot := strings.Index(fn, "."); dot >= 0 {
		pkg += fn[:dot]
		fn = fn[dot+1:]
	}

	fn = strings.Replace(fn, "·", ".", -1)
	return
}

// Source: https://github.com/golang/go/blob/master/src/cmd/go/test.go
// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
	return !unicode.IsLower(rune)
}

func getWhitespaceString() string {

	_, file, line, ok := runtime.Caller(1)
	if !ok {
		return ""
	}
	parts := strings.Split(file, "/")
	file = parts[len(parts)-1]

	return strings.Repeat(" ", len(fmt.Sprintf("%s:%d:      ", file, line)))

}