/
servertest.go
126 lines (109 loc) · 3.36 KB
/
servertest.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
// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package servertest provides utilities for running tests against a remote LSP
// server.
package servertest
import (
"context"
"fmt"
"io"
"net"
"sync"
"github.com/prometheus-community/promql-langserver/internal/vendored/go-tools/jsonrpc2"
)
// Connector is the interface used to connect to a server.
type Connector interface {
Connect(context.Context) *jsonrpc2.Conn
}
// TCPServer is a helper for executing tests against a remote jsonrpc2
// connection. Once initialized, its Addr field may be used to connect a
// jsonrpc2 client.
type TCPServer struct {
Addr string
ln net.Listener
cls *closerList
}
// NewTCPServer returns a new test server listening on local tcp port and
// serving incoming jsonrpc2 streams using the provided stream server. It
// panics on any error.
func NewTCPServer(ctx context.Context, server jsonrpc2.StreamServer) *TCPServer {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
panic(fmt.Sprintf("servertest: failed to listen: %v", err))
}
go jsonrpc2.Serve(ctx, ln, server, 0)
return &TCPServer{Addr: ln.Addr().String(), ln: ln, cls: &closerList{}}
}
// Connect dials the test server and returns a jsonrpc2 Connection that is
// ready for use.
func (s *TCPServer) Connect(ctx context.Context) *jsonrpc2.Conn {
netConn, err := net.Dial("tcp", s.Addr)
if err != nil {
panic(fmt.Sprintf("servertest: failed to connect to test instance: %v", err))
}
s.cls.add(func() {
netConn.Close()
})
conn := jsonrpc2.NewConn(jsonrpc2.NewHeaderStream(netConn, netConn))
go conn.Run(ctx)
return conn
}
// Close closes all connected pipes.
func (s *TCPServer) Close() error {
s.cls.closeAll()
return nil
}
// PipeServer is a test server that handles connections over io.Pipes.
type PipeServer struct {
server jsonrpc2.StreamServer
cls *closerList
}
// NewPipeServer returns a test server that can be connected to via io.Pipes.
func NewPipeServer(ctx context.Context, server jsonrpc2.StreamServer) *PipeServer {
return &PipeServer{server: server, cls: &closerList{}}
}
// Connect creates new io.Pipes and binds them to the underlying StreamServer.
func (s *PipeServer) Connect(ctx context.Context) *jsonrpc2.Conn {
// Pipes connect like this:
// Client🡒(sWriter)🡒(sReader)🡒Server
// 🡔(cReader)🡐(cWriter)🡗
sReader, sWriter := io.Pipe()
cReader, cWriter := io.Pipe()
s.cls.add(func() {
sReader.Close()
sWriter.Close()
cReader.Close()
cWriter.Close()
})
serverStream := jsonrpc2.NewStream(sReader, cWriter)
go s.server.ServeStream(ctx, serverStream)
clientStream := jsonrpc2.NewStream(cReader, sWriter)
clientConn := jsonrpc2.NewConn(clientStream)
go clientConn.Run(ctx)
return clientConn
}
// Close closes all connected pipes.
func (s *PipeServer) Close() error {
s.cls.closeAll()
return nil
}
// closerList tracks closers to run when a testserver is closed. This is a
// convenience, so that callers don't have to worry about closing each
// connection.
type closerList struct {
mu sync.Mutex
closers []func()
}
func (l *closerList) add(closer func()) {
l.mu.Lock()
defer l.mu.Unlock()
l.closers = append(l.closers, closer)
}
func (l *closerList) closeAll() {
l.mu.Lock()
defer l.mu.Unlock()
for _, closer := range l.closers {
closer()
}
}