forked from cloudfoundry-attic/app-manager
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_helpers.go
85 lines (66 loc) · 1.74 KB
/
test_helpers.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 test_helpers
import (
"errors"
"os"
"sync"
"github.com/tedsuo/ifrit"
)
// PingChan stops when you send it a single Ping
type PingChan chan Ping
type Ping struct{}
var PingerExitedFromPing = errors.New("pinger exited with a ping")
var PingerExitedFromSignal = errors.New("pinger exited with a signal")
func (p PingChan) Load(err error) (ifrit.Runner, bool) {
return p, true
}
func (p PingChan) Run(sigChan <-chan os.Signal, ready chan<- struct{}) error {
close(ready)
select {
case <-sigChan:
return PingerExitedFromSignal
case p <- Ping{}:
return PingerExitedFromPing
}
}
// NoReadyRunner exits without closing the ready chan
var NoReadyRunner = ifrit.RunFunc(func(sigChan <-chan os.Signal, ready chan<- struct{}) error {
return NoReadyExitedNormally
})
var NoReadyExitedNormally = errors.New("no ready exited normally")
// SignalRecoder records all signals received, and exits on a set of signals.
type SignalRecoder struct {
sync.RWMutex
signals []os.Signal
exitSignals map[os.Signal]struct{}
}
func NewSignalRecorder(exitSignals ...os.Signal) *SignalRecoder {
exitSignals = append(exitSignals, os.Kill, os.Interrupt)
signalSet := map[os.Signal]struct{}{}
for _, signal := range exitSignals {
signalSet[signal] = struct{}{}
}
return &SignalRecoder{
exitSignals: signalSet,
}
}
func (r *SignalRecoder) Load(err error) (ifrit.Runner, bool) {
return r, true
}
func (r *SignalRecoder) Run(sigChan <-chan os.Signal, ready chan<- struct{}) error {
close(ready)
for {
signal := <-sigChan
r.Lock()
r.signals = append(r.signals, signal)
r.Unlock()
_, ok := r.exitSignals[signal]
if ok {
return nil
}
}
}
func (r *SignalRecoder) ReceivedSignals() []os.Signal {
defer r.RUnlock()
r.RLock()
return r.signals
}