/
nettest.go
50 lines (46 loc) · 1.12 KB
/
nettest.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
/*
Package nettest provides a test server to test network listener handlers,
similar to how net/http/httptest provides HTTP handler testing.
*/
package nettest
import (
"context"
"crypto/tls"
"fmt"
"net"
)
type Server struct {
Addr string
lis net.Listener
cancel func()
done chan struct{}
}
func (s *Server) Close() error {
defer func() { <-s.done }()
s.cancel()
return s.lis.Close()
}
// NewTLSServer creates a TLS TCP server listening on 127.0.0.1:0 and runs
// handler with the listener.
// The listen address with the chosen port can be read from the server's Addr field.
// Clients should tls.Dial using ClientTLS as the TLS config.
func NewTLSServer(handler func(lis net.Listener)) *Server {
ctx, cancel := context.WithCancel(context.Background())
var lc net.ListenConfig
lis, err := lc.Listen(ctx, "tcp", "127.0.0.1:0")
if err != nil {
panic(fmt.Sprintf("nettest: NewTLSServer: %v", err))
}
lis = tls.NewListener(lis, serverTLS)
svr := &Server{
Addr: lis.Addr().String(),
lis: lis,
cancel: cancel,
done: make(chan struct{}),
}
go func() {
handler(lis)
close(svr.done)
}()
return svr
}