github.com/decred/dcrlnd@v0.7.6/clock/test_clock.go (about)

     1  package clock
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  )
     7  
     8  // TestClock can be used in tests to mock time.
     9  type TestClock struct {
    10  	currentTime time.Time
    11  	timeChanMap map[time.Time][]chan time.Time
    12  	timeLock    sync.Mutex
    13  	tickSignal  chan time.Duration
    14  }
    15  
    16  // NewTestClock returns a new test clock.
    17  func NewTestClock(startTime time.Time) *TestClock {
    18  	return &TestClock{
    19  		currentTime: startTime,
    20  		timeChanMap: make(map[time.Time][]chan time.Time),
    21  	}
    22  }
    23  
    24  // NewTestClockWithTickSignal will create a new test clock with an added
    25  // channel which will be used to signal when a new ticker is registered.
    26  // This is useful when creating a ticker on a separate goroutine and we'd
    27  // like to wait for that to happen before advancing the test case.
    28  func NewTestClockWithTickSignal(startTime time.Time,
    29  	tickSignal chan time.Duration) *TestClock {
    30  
    31  	testClock := NewTestClock(startTime)
    32  	testClock.tickSignal = tickSignal
    33  
    34  	return testClock
    35  }
    36  
    37  // Now returns the current (test) time.
    38  func (c *TestClock) Now() time.Time {
    39  	c.timeLock.Lock()
    40  	defer c.timeLock.Unlock()
    41  
    42  	return c.currentTime
    43  }
    44  
    45  // TickAfter returns a channel that will receive a tick after the specified
    46  // duration has passed passed by the user set test time.
    47  func (c *TestClock) TickAfter(duration time.Duration) <-chan time.Time {
    48  	c.timeLock.Lock()
    49  	defer func() {
    50  		c.timeLock.Unlock()
    51  
    52  		// Signal that the ticker has been added.
    53  		if c.tickSignal != nil {
    54  			c.tickSignal <- duration
    55  		}
    56  	}()
    57  
    58  	triggerTime := c.currentTime.Add(duration)
    59  	ch := make(chan time.Time, 1)
    60  
    61  	// If already expired, tick immediately.
    62  	if !triggerTime.After(c.currentTime) {
    63  		ch <- c.currentTime
    64  		return ch
    65  	}
    66  
    67  	// Otherwise store the channel until the trigger time is there.
    68  	chans := c.timeChanMap[triggerTime]
    69  	chans = append(chans, ch)
    70  	c.timeChanMap[triggerTime] = chans
    71  
    72  	return ch
    73  }
    74  
    75  // SetTime sets the (test) time and triggers tick channels when they expire.
    76  func (c *TestClock) SetTime(now time.Time) {
    77  	c.timeLock.Lock()
    78  	defer c.timeLock.Unlock()
    79  
    80  	c.currentTime = now
    81  	remainingChans := make(map[time.Time][]chan time.Time)
    82  	for triggerTime, chans := range c.timeChanMap {
    83  		// If the trigger time is still in the future, keep this channel
    84  		// in the channel map for later.
    85  		if triggerTime.After(now) {
    86  			remainingChans[triggerTime] = chans
    87  			continue
    88  		}
    89  
    90  		for _, c := range chans {
    91  			c <- now
    92  		}
    93  	}
    94  
    95  	c.timeChanMap = remainingChans
    96  }