code_location.go

Documentation: github.com/onsi/ginkgo/v2/types

     1  package types
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"regexp"
     7  	"runtime"
     8  	"runtime/debug"
     9  	"strings"
    10  	"sync"
    11  )
    12  
    13  type CodeLocation struct {
    14  	FileName       string `json:",omitempty"`
    15  	LineNumber     int    `json:",omitempty"`
    16  	FullStackTrace string `json:",omitempty"`
    17  	CustomMessage  string `json:",omitempty"`
    18  }
    19  
    20  func (codeLocation CodeLocation) String() string {
    21  	if codeLocation.CustomMessage != "" {
    22  		return codeLocation.CustomMessage
    23  	}
    24  	return fmt.Sprintf("%s:%d", codeLocation.FileName, codeLocation.LineNumber)
    25  }
    26  
    27  func (codeLocation CodeLocation) ContentsOfLine() string {
    28  	if codeLocation.CustomMessage != "" {
    29  		return ""
    30  	}
    31  	contents, err := os.ReadFile(codeLocation.FileName)
    32  	if err != nil {
    33  		return ""
    34  	}
    35  	lines := strings.Split(string(contents), "\n")
    36  	if len(lines) < codeLocation.LineNumber {
    37  		return ""
    38  	}
    39  	return lines[codeLocation.LineNumber-1]
    40  }
    41  
    42  type codeLocationLocator struct {
    43  	pcs     map[uintptr]bool
    44  	helpers map[string]bool
    45  	lock    *sync.Mutex
    46  }
    47  
    48  func (c *codeLocationLocator) addHelper(pc uintptr) {
    49  	c.lock.Lock()
    50  	defer c.lock.Unlock()
    51  
    52  	if c.pcs[pc] {
    53  		return
    54  	}
    55  	c.lock.Unlock()
    56  	f := runtime.FuncForPC(pc)
    57  	c.lock.Lock()
    58  	if f == nil {
    59  		return
    60  	}
    61  	c.helpers[f.Name()] = true
    62  	c.pcs[pc] = true
    63  }
    64  
    65  func (c *codeLocationLocator) hasHelper(name string) bool {
    66  	c.lock.Lock()
    67  	defer c.lock.Unlock()
    68  	return c.helpers[name]
    69  }
    70  
    71  func (c *codeLocationLocator) getCodeLocation(skip int) CodeLocation {
    72  	pc := make([]uintptr, 40)
    73  	n := runtime.Callers(skip+2, pc)
    74  	if n == 0 {
    75  		return CodeLocation{}
    76  	}
    77  	pc = pc[:n]
    78  	frames := runtime.CallersFrames(pc)
    79  	for {
    80  		frame, more := frames.Next()
    81  		if !c.hasHelper(frame.Function) {
    82  			return CodeLocation{FileName: frame.File, LineNumber: frame.Line}
    83  		}
    84  		if !more {
    85  			break
    86  		}
    87  	}
    88  	return CodeLocation{}
    89  }
    90  
    91  var clLocator = &codeLocationLocator{
    92  	pcs:     map[uintptr]bool{},
    93  	helpers: map[string]bool{},
    94  	lock:    &sync.Mutex{},
    95  }
    96  
    97  // MarkAsHelper is used by GinkgoHelper to mark the caller (appropriately offset by skip)as a helper.  You can use this directly if you need to provide an optional `skip` to mark functions further up the call stack as helpers.
    98  func MarkAsHelper(optionalSkip ...int) {
    99  	skip := 1
   100  	if len(optionalSkip) > 0 {
   101  		skip += optionalSkip[0]
   102  	}
   103  	pc, _, _, ok := runtime.Caller(skip)
   104  	if ok {
   105  		clLocator.addHelper(pc)
   106  	}
   107  }
   108  
   109  func NewCustomCodeLocation(message string) CodeLocation {
   110  	return CodeLocation{
   111  		CustomMessage: message,
   112  	}
   113  }
   114  
   115  func NewCodeLocation(skip int) CodeLocation {
   116  	return clLocator.getCodeLocation(skip + 1)
   117  }
   118  
   119  func NewCodeLocationWithStackTrace(skip int) CodeLocation {
   120  	cl := clLocator.getCodeLocation(skip + 1)
   121  	cl.FullStackTrace = PruneStack(string(debug.Stack()), skip+1)
   122  	return cl
   123  }
   124  
   125  // PruneStack removes references to functions that are internal to Ginkgo
   126  // and the Go runtime from a stack string and a certain number of stack entries
   127  // at the beginning of the stack. The stack string has the format
   128  // as returned by runtime/debug.Stack. The leading goroutine information is
   129  // optional and always removed if present. Beware that runtime/debug.Stack
   130  // adds itself as first entry, so typically skip must be >= 1 to remove that
   131  // entry.
   132  func PruneStack(fullStackTrace string, skip int) string {
   133  	stack := strings.Split(fullStackTrace, "\n")
   134  	// Ensure that the even entries are the method names and the
   135  	// odd entries the source code information.
   136  	if len(stack) > 0 && strings.HasPrefix(stack[0], "goroutine ") {
   137  		// Ignore "goroutine 29 [running]:" line.
   138  		stack = stack[1:]
   139  	}
   140  	// The "+1" is for skipping over the initial entry, which is
   141  	// runtime/debug.Stack() itself.
   142  	if len(stack) > 2*(skip+1) {
   143  		stack = stack[2*(skip+1):]
   144  	}
   145  	prunedStack := []string{}
   146  	if os.Getenv("GINKGO_PRUNE_STACK") == "FALSE" {
   147  		prunedStack = stack
   148  	} else {
   149  		re := regexp.MustCompile(`\/ginkgo\/|\/pkg\/testing\/|\/pkg\/runtime\/`)
   150  		for i := 0; i < len(stack)/2; i++ {
   151  			// We filter out based on the source code file name.
   152  			if !re.MatchString(stack[i*2+1]) {
   153  				prunedStack = append(prunedStack, stack[i*2])
   154  				prunedStack = append(prunedStack, stack[i*2+1])
   155  			}
   156  		}
   157  	}
   158  	return strings.Join(prunedStack, "\n")
   159  }
   160
View as plain text