github.com/v2fly/tools@v0.100.0/internal/jsonrpc2/servertest/servertest.go (about)

     1  // Copyright 2020 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package servertest provides utilities for running tests against a remote LSP
     6  // server.
     7  package servertest
     8  
     9  import (
    10  	"context"
    11  	"fmt"
    12  	"net"
    13  	"strings"
    14  	"sync"
    15  
    16  	"github.com/v2fly/tools/internal/jsonrpc2"
    17  )
    18  
    19  // Connector is the interface used to connect to a server.
    20  type Connector interface {
    21  	Connect(context.Context) jsonrpc2.Conn
    22  }
    23  
    24  // TCPServer is a helper for executing tests against a remote jsonrpc2
    25  // connection. Once initialized, its Addr field may be used to connect a
    26  // jsonrpc2 client.
    27  type TCPServer struct {
    28  	*connList
    29  
    30  	Addr string
    31  
    32  	ln     net.Listener
    33  	framer jsonrpc2.Framer
    34  }
    35  
    36  // NewTCPServer returns a new test server listening on local tcp port and
    37  // serving incoming jsonrpc2 streams using the provided stream server. It
    38  // panics on any error.
    39  func NewTCPServer(ctx context.Context, server jsonrpc2.StreamServer, framer jsonrpc2.Framer) *TCPServer {
    40  	ln, err := net.Listen("tcp", "127.0.0.1:0")
    41  	if err != nil {
    42  		panic(fmt.Sprintf("servertest: failed to listen: %v", err))
    43  	}
    44  	if framer == nil {
    45  		framer = jsonrpc2.NewHeaderStream
    46  	}
    47  	go jsonrpc2.Serve(ctx, ln, server, 0)
    48  	return &TCPServer{Addr: ln.Addr().String(), ln: ln, framer: framer, connList: &connList{}}
    49  }
    50  
    51  // Connect dials the test server and returns a jsonrpc2 Connection that is
    52  // ready for use.
    53  func (s *TCPServer) Connect(ctx context.Context) jsonrpc2.Conn {
    54  	netConn, err := net.Dial("tcp", s.Addr)
    55  	if err != nil {
    56  		panic(fmt.Sprintf("servertest: failed to connect to test instance: %v", err))
    57  	}
    58  	conn := jsonrpc2.NewConn(s.framer(netConn))
    59  	s.add(conn)
    60  	return conn
    61  }
    62  
    63  // PipeServer is a test server that handles connections over io.Pipes.
    64  type PipeServer struct {
    65  	*connList
    66  	server jsonrpc2.StreamServer
    67  	framer jsonrpc2.Framer
    68  }
    69  
    70  // NewPipeServer returns a test server that can be connected to via io.Pipes.
    71  func NewPipeServer(ctx context.Context, server jsonrpc2.StreamServer, framer jsonrpc2.Framer) *PipeServer {
    72  	if framer == nil {
    73  		framer = jsonrpc2.NewRawStream
    74  	}
    75  	return &PipeServer{server: server, framer: framer, connList: &connList{}}
    76  }
    77  
    78  // Connect creates new io.Pipes and binds them to the underlying StreamServer.
    79  func (s *PipeServer) Connect(ctx context.Context) jsonrpc2.Conn {
    80  	sPipe, cPipe := net.Pipe()
    81  	serverStream := s.framer(sPipe)
    82  	serverConn := jsonrpc2.NewConn(serverStream)
    83  	s.add(serverConn)
    84  	go s.server.ServeStream(ctx, serverConn)
    85  
    86  	clientStream := s.framer(cPipe)
    87  	clientConn := jsonrpc2.NewConn(clientStream)
    88  	s.add(clientConn)
    89  	return clientConn
    90  }
    91  
    92  // connList tracks closers to run when a testserver is closed.  This is a
    93  // convenience, so that callers don't have to worry about closing each
    94  // connection.
    95  type connList struct {
    96  	mu    sync.Mutex
    97  	conns []jsonrpc2.Conn
    98  }
    99  
   100  func (l *connList) add(conn jsonrpc2.Conn) {
   101  	l.mu.Lock()
   102  	defer l.mu.Unlock()
   103  	l.conns = append(l.conns, conn)
   104  }
   105  
   106  func (l *connList) Close() error {
   107  	l.mu.Lock()
   108  	defer l.mu.Unlock()
   109  	var errmsgs []string
   110  	for _, conn := range l.conns {
   111  		if err := conn.Close(); err != nil {
   112  			errmsgs = append(errmsgs, err.Error())
   113  		}
   114  	}
   115  	if len(errmsgs) > 0 {
   116  		return fmt.Errorf("closing errors:\n%s", strings.Join(errmsgs, "\n"))
   117  	}
   118  	return nil
   119  }