/
brain.go
85 lines (70 loc) · 1.92 KB
/
brain.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package joetest
import (
"context"
"sync"
"time"
"github.com/go-joe/joe"
"go.uber.org/zap/zaptest"
)
// Brain wraps the joe.Brain for unit testing.
type Brain struct {
*joe.Brain
mu sync.Mutex
events []interface{}
eventsChan chan joe.Event
}
// NewBrain creates a new Brain that can be used for unit testing. The Brain
// registers to all events except the (init and shutdown event) and records them
// for later access. The event handling loop of the Brain (i.e. Brain.HandleEvents())
// is automatically started by this function in a new goroutine and the caller
// must call Brain.Finish() at the end of their tests.
func NewBrain(t TestingT) *Brain {
logger := zaptest.NewLogger(t)
b := &Brain{
Brain: joe.NewBrain(logger),
eventsChan: make(chan joe.Event, 100),
}
initialized := make(chan bool)
b.RegisterHandler(b.observeEvent)
b.RegisterHandler(func(joe.InitEvent) {
initialized <- true
})
go b.HandleEvents()
<-initialized
return b
}
func (b *Brain) observeEvent(evt interface{}) {
switch evt.(type) {
case joe.InitEvent, joe.ShutdownEvent:
return
}
select {
case b.eventsChan <- joe.Event{Data: evt}:
// ok, lets move on
default:
// nobody is listening, also fine
}
b.mu.Lock()
b.events = append(b.events, evt)
b.mu.Unlock()
}
// Finish stops the event handler loop of the Brain and waits until all pending
// events have been processed.
func (b *Brain) Finish() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
b.Brain.Shutdown(ctx)
}
// RecordedEvents returns all events the Brain has processed except the
// joe.InitEvent and joe.ShutdownEvent.
func (b *Brain) RecordedEvents() []interface{} {
b.mu.Lock()
events := make([]interface{}, len(b.events))
copy(events, b.events)
b.mu.Unlock()
return events
}
// Events returns a channel that receives all emitted events.
func (b *Brain) Events() <-chan joe.Event {
return b.eventsChan
}