-
Notifications
You must be signed in to change notification settings - Fork 16
/
mock_stream.go
75 lines (67 loc) · 1.35 KB
/
mock_stream.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
package mock
import (
"context"
"fmt"
"sync"
"sync/atomic"
"testing"
)
type MockStream[T any] struct {
mu sync.Mutex
ctx context.Context
ch chan *T
closed bool
counter *uint32
Messages []*T
MessageMap map[string]int
}
func NewMockStream[T any](t *testing.T) *MockStream[T] {
t.Helper()
var counter uint32
return &MockStream[T]{
ctx: context.Background(),
ch: make(chan *T),
closed: false,
counter: &counter,
Messages: make([]*T, 0),
MessageMap: make(map[string]int),
}
}
func (m *MockStream[T]) Run() error {
for {
select {
case data, ok := <-m.ch:
if !ok {
return fmt.Errorf("stream closed")
}
atomic.AddUint32(m.counter, 1)
m.Messages = append(m.Messages, data)
case <-m.ctx.Done():
return m.ctx.Err()
}
}
}
func (m *MockStream[T]) GetChannel() chan *T {
return m.ch
}
func (m *MockStream[T]) Send(data *T) {
m.mu.Lock()
defer m.mu.Unlock()
if !m.closed {
m.ch <- data
}
}
func (m *MockStream[T]) Close() {
m.mu.Lock()
defer m.mu.Unlock()
if !m.closed {
close(m.ch)
}
m.closed = true
}
// SetCounter sets the counter used to count the number of messages sent.
// Multiple streams can share the same counter to count the total number of
// messages sent across all streams.
func (m *MockStream[T]) SetCounter(counter *uint32) {
m.counter = counter
}