145 lines
2.6 KiB
Go
145 lines
2.6 KiB
Go
package elect_test
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/dchest/uniuri"
|
|
"github.com/gopatchy/elect"
|
|
"github.com/gopatchy/event"
|
|
"github.com/gopatchy/proxy"
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type TestServer struct {
|
|
Candidate *elect.Candidate
|
|
|
|
listener *net.TCPListener
|
|
srv *http.Server
|
|
}
|
|
|
|
type TestSystem struct {
|
|
signingKey string
|
|
servers []*TestServer
|
|
voters []*elect.Voter
|
|
proxies []*proxy.Proxy
|
|
}
|
|
|
|
type Waiter struct {
|
|
chans []<-chan bool
|
|
}
|
|
|
|
func NewTestServer(ctx context.Context, t *testing.T, ec *event.Client, numVoters int, signingKey string) *TestServer {
|
|
ts := &TestServer{
|
|
Candidate: elect.NewCandidate(ctx, ec, numVoters, signingKey),
|
|
listener: lo.Must(net.ListenTCP("tcp", nil)),
|
|
}
|
|
|
|
ts.srv = &http.Server{
|
|
Handler: ts.Candidate,
|
|
ReadHeaderTimeout: 30 * time.Second,
|
|
}
|
|
|
|
go func() {
|
|
err := ts.srv.Serve(ts.listener)
|
|
require.ErrorIs(t, err, http.ErrServerClosed)
|
|
}()
|
|
|
|
return ts
|
|
}
|
|
|
|
func (ts *TestServer) Stop() {
|
|
ts.srv.Close()
|
|
ts.Candidate.Stop()
|
|
}
|
|
|
|
func (ts *TestServer) Addr() *net.TCPAddr {
|
|
return ts.listener.Addr().(*net.TCPAddr)
|
|
}
|
|
|
|
func NewTestSystem(t *testing.T, numCandidates, numVoters int) *TestSystem {
|
|
ctx := context.Background()
|
|
|
|
ec := event.New()
|
|
defer ec.Close()
|
|
|
|
ts := &TestSystem{
|
|
signingKey: uniuri.New(),
|
|
}
|
|
|
|
for i := 0; i < numCandidates; i++ {
|
|
ts.servers = append(ts.servers, NewTestServer(ctx, t, ec, numVoters, ts.signingKey))
|
|
}
|
|
|
|
for i := 0; i < numVoters; i++ {
|
|
ts.proxies = append(ts.proxies, proxy.NewProxy(t, ts.Server(0).Addr()))
|
|
ts.voters = append(ts.voters, elect.NewVoter(ctx, ec, ts.Proxy(i).HTTP(), ts.signingKey))
|
|
}
|
|
|
|
return ts
|
|
}
|
|
|
|
func (ts *TestSystem) Stop() {
|
|
for _, s := range ts.servers {
|
|
s.Stop()
|
|
}
|
|
|
|
for _, v := range ts.voters {
|
|
v.Stop()
|
|
}
|
|
|
|
for _, p := range ts.proxies {
|
|
p.Close()
|
|
}
|
|
}
|
|
|
|
func (ts *TestSystem) SetServer(server int) {
|
|
for _, p := range ts.proxies {
|
|
p.SetBackend(ts.Server(server).Addr())
|
|
}
|
|
}
|
|
|
|
func (ts *TestSystem) SetServerForVoter(server, voter int) {
|
|
ts.Proxy(voter).SetBackend(ts.Server(server).Addr())
|
|
}
|
|
|
|
func (ts *TestSystem) Candidate(i int) *elect.Candidate {
|
|
return ts.servers[i].Candidate
|
|
}
|
|
|
|
func (ts *TestSystem) Proxy(i int) *proxy.Proxy {
|
|
return ts.proxies[i]
|
|
}
|
|
|
|
func (ts *TestSystem) Server(i int) *TestServer {
|
|
return ts.servers[i]
|
|
}
|
|
|
|
func (ts *TestSystem) Voter(i int) *elect.Voter {
|
|
return ts.voters[i]
|
|
}
|
|
|
|
func NewWaiter() *Waiter {
|
|
return &Waiter{}
|
|
}
|
|
|
|
func (w *Waiter) Wait() {
|
|
for _, ch := range w.chans {
|
|
<-ch
|
|
}
|
|
}
|
|
|
|
func (w *Waiter) Async(cb func()) {
|
|
ch := make(chan bool)
|
|
w.chans = append(w.chans, ch)
|
|
|
|
go func() {
|
|
defer close(ch)
|
|
cb()
|
|
}()
|
|
}
|