/
mock_connection.go
112 lines (93 loc) · 2.07 KB
/
mock_connection.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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
package gostratum
import (
"context"
"fmt"
"net"
"sync"
"sync/atomic"
"time"
)
type MockConnection struct {
id string
lock sync.Mutex // to prevent double closing of channel
inChan chan []byte
outChan chan []byte
}
var channelCounter int32
func NewMockConnection() *MockConnection {
return &MockConnection{
id: fmt.Sprintf("mc_%d", atomic.AddInt32(&channelCounter, 1)),
lock: sync.Mutex{},
inChan: make(chan []byte),
outChan: make(chan []byte),
}
}
func (mc *MockConnection) AsyncWriteTestDataToReadBuffer(s string) {
go func() {
mc.inChan <- []byte(s)
}()
}
func (mc *MockConnection) ReadTestDataFromBuffer(handler func([]byte)) {
read := <-mc.outChan
handler(read)
}
func (mc *MockConnection) AsyncReadTestDataFromBuffer(handler func([]byte)) {
go func() {
read := <-mc.outChan
handler(read)
}()
}
func (mc *MockConnection) Read(b []byte) (int, error) {
data, ok := <-mc.inChan
if !ok {
return 0, context.DeadlineExceeded
}
return copy(b, data), nil
}
func (mc *MockConnection) Write(b []byte) (int, error) {
mc.outChan <- b
return len(b), nil
}
func (mc *MockConnection) Close() error {
mc.lock.Lock()
defer mc.lock.Unlock()
close(mc.inChan)
close(mc.outChan)
return nil
}
type MockAddr struct {
id string
}
func (ma MockAddr) Network() string { return "mock" }
func (ma MockAddr) String() string { return ma.id }
func (mc *MockConnection) LocalAddr() net.Addr {
return MockAddr{id: mc.id}
}
func (mc *MockConnection) RemoteAddr() net.Addr {
return MockAddr{id: mc.id}
}
func (mc *MockConnection) SetDeadline(t time.Time) error {
mc.SetReadDeadline(t)
mc.SetWriteDeadline(t)
return nil
}
func (mc *MockConnection) SetReadDeadline(t time.Time) error {
go func() {
mc.lock.Lock()
defer mc.lock.Unlock()
time.Sleep(time.Until(t))
close(mc.inChan)
mc.inChan = make(chan []byte)
}()
return nil
}
func (mc *MockConnection) SetWriteDeadline(t time.Time) error {
go func() {
mc.lock.Lock()
defer mc.lock.Unlock()
time.Sleep(time.Until(t))
close(mc.outChan)
mc.outChan = make(chan []byte)
}()
return nil
}