github.com/tenntenn/testtime@v0.2.3-0.20221118081726-55bcd1f05226/time.go (about)

     1  package testtime
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  	_ "unsafe" // for go:linkname
     9  )
    10  
    11  //go:linkname timeMap time.timeMap
    12  var timeMap sync.Map
    13  
    14  // SetTime stores a fixed time which can be identified by the caller's goroutine.
    15  // The fixed time will be deleted at the end of test.
    16  func SetTime(t *testing.T, tm time.Time) {
    17  	t.Helper()
    18  
    19  	key := goroutineID()
    20  	timeMap.Store(key, func() time.Time {
    21  		return tm
    22  	})
    23  
    24  	t.Cleanup(func() {
    25  		timeMap.Delete(key)
    26  	})
    27  }
    28  
    29  // SetFunc stores a function which returns time.Time which can be identified by the caller's goroutine.
    30  // The function will be deleted at the end of test.
    31  func SetFunc(t *testing.T, f func() time.Time) {
    32  	t.Helper()
    33  
    34  	key := goroutineID()
    35  	timeMap.Store(key, f)
    36  
    37  	t.Cleanup(func() {
    38  		timeMap.Delete(key)
    39  	})
    40  }
    41  
    42  // Now returns a fixed time which is related with the goroutine by SetTime or SetFunc.
    43  // If the current goroutine is not related with any fixed time or function, Now calls time.Now and returns its returned value.
    44  func Now() time.Time {
    45  	v, ok := timeMap.Load(goroutineID())
    46  	if ok {
    47  		return v.(func() time.Time)()
    48  	}
    49  	return time.Now()
    50  }
    51  
    52  func goroutineID() string {
    53  	var buf [64]byte
    54  	n := runtime.Stack(buf[:], false)
    55  	// 10: len("goroutine ")
    56  	for i := 10; i < n; i++ {
    57  		if buf[i] == ' ' {
    58  			return string(buf[10:i])
    59  		}
    60  	}
    61  	return ""
    62  }