-
Notifications
You must be signed in to change notification settings - Fork 53
/
haproxy_mock.go
135 lines (118 loc) · 2.55 KB
/
haproxy_mock.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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
package runtime
import (
"bufio"
"errors"
"io"
"net"
"os"
"strings"
"sync"
"testing"
)
// HAProxyMock - Mock HAProxy Server for testing the socket communication
type HAProxyMock struct {
net.Listener
responses map[string]string
t *testing.T
mu sync.RWMutex
running bool
}
// NewHAProxyMock - create new haproxy mock
func NewHAProxyMock(t *testing.T) *HAProxyMock {
haProxyMock := &HAProxyMock{}
haProxyMock.t = t
l, err := net.Listen("unix", socket())
if err != nil {
t.Fatal(err)
}
haProxyMock.Listener = l
return haProxyMock
}
// Stop - stop mock
func (haproxy *HAProxyMock) Stop() {
haproxy.setRunning(false)
}
// Start - start mock
func (haproxy *HAProxyMock) Start() {
haproxy.setRunning(true)
go func() {
for {
if !haproxy.getRunning() {
return
}
conn, err := haproxy.Accept()
if err != nil {
haproxy.t.Error(err)
return
}
haproxy.handleConnection(conn)
}
}()
}
// SetResponses - setting the expected responses, safe for concurrent use
func (haproxy *HAProxyMock) SetResponses(responses *map[string]string) {
haproxy.mu.Lock()
haproxy.responses = *responses
haproxy.mu.Unlock()
}
// getResponses gets the responses of the mock, safe for concurrent use
func (haproxy *HAProxyMock) getResponses() map[string]string {
haproxy.mu.RLock()
defer haproxy.mu.RUnlock()
return haproxy.responses
}
// setRunning sets the running state of the mock, safe for concurrent use
func (haproxy *HAProxyMock) setRunning(running bool) {
haproxy.mu.Lock()
haproxy.running = running
haproxy.mu.Unlock()
}
// getRunning gets the running state of the mock, safe for concurrent use
func (haproxy *HAProxyMock) getRunning() bool {
haproxy.mu.RLock()
defer haproxy.mu.RUnlock()
return haproxy.running
}
func (haproxy *HAProxyMock) handleConnection(conn net.Conn) {
go func() {
defer func() {
_ = conn.Close()
}()
r := bufio.NewReader(conn)
w := bufio.NewWriter(conn)
s, err := r.ReadString('\n')
if err != nil && errors.Is(err, io.EOF) {
haproxy.t.Log(err)
return
}
if strings.Contains(s, "<<") {
r, _ := r.ReadString('\n')
s += r
}
split := strings.Split(s, ";")
if len(split) > 0 {
s = split[len(split)-1]
}
response := haproxy.getResponses()[s]
_, err = w.WriteString(response)
if err != nil {
haproxy.t.Log(err)
return
}
err = w.Flush()
if err != nil {
haproxy.t.Log(err)
return
}
}()
}
func socket() string {
f, err := os.CreateTemp("", "haproxy-sock")
if err != nil {
panic(err)
}
addr := f.Name()
_ = f.Close()
_ = os.Remove(addr)
return addr
}