github.com/v2fly/tools@v0.100.0/internal/jsonrpc2_v2/serve_test.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 jsonrpc2_test
     6  
     7  import (
     8  	"context"
     9  	"errors"
    10  	"testing"
    11  	"time"
    12  
    13  	jsonrpc2 "github.com/v2fly/tools/internal/jsonrpc2_v2"
    14  	"github.com/v2fly/tools/internal/stack/stacktest"
    15  )
    16  
    17  func TestIdleTimeout(t *testing.T) {
    18  	stacktest.NoLeak(t)
    19  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
    20  	defer cancel()
    21  
    22  	listener, err := jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
    23  	if err != nil {
    24  		t.Fatal(err)
    25  	}
    26  	listener = jsonrpc2.NewIdleListener(100*time.Millisecond, listener)
    27  	defer listener.Close()
    28  	server, err := jsonrpc2.Serve(ctx, listener, jsonrpc2.ConnectionOptions{})
    29  	if err != nil {
    30  		t.Fatal(err)
    31  	}
    32  
    33  	connect := func() *jsonrpc2.Connection {
    34  		client, err := jsonrpc2.Dial(ctx,
    35  			listener.Dialer(),
    36  			jsonrpc2.ConnectionOptions{})
    37  		if err != nil {
    38  			t.Fatal(err)
    39  		}
    40  		return client
    41  	}
    42  	// Exercise some connection/disconnection patterns, and then assert that when
    43  	// our timer fires, the server exits.
    44  	conn1 := connect()
    45  	conn2 := connect()
    46  	if err := conn1.Close(); err != nil {
    47  		t.Fatalf("conn1.Close failed with error: %v", err)
    48  	}
    49  	if err := conn2.Close(); err != nil {
    50  		t.Fatalf("conn2.Close failed with error: %v", err)
    51  	}
    52  	conn3 := connect()
    53  	if err := conn3.Close(); err != nil {
    54  		t.Fatalf("conn3.Close failed with error: %v", err)
    55  	}
    56  
    57  	serverError := server.Wait()
    58  
    59  	if !errors.Is(serverError, jsonrpc2.ErrIdleTimeout) {
    60  		t.Errorf("run() returned error %v, want %v", serverError, jsonrpc2.ErrIdleTimeout)
    61  	}
    62  }
    63  
    64  type msg struct {
    65  	Msg string
    66  }
    67  
    68  type fakeHandler struct{}
    69  
    70  func (fakeHandler) Handle(ctx context.Context, req *jsonrpc2.Request) (interface{}, error) {
    71  	switch req.Method {
    72  	case "ping":
    73  		return &msg{"pong"}, nil
    74  	default:
    75  		return nil, jsonrpc2.ErrNotHandled
    76  	}
    77  }
    78  
    79  func TestServe(t *testing.T) {
    80  	stacktest.NoLeak(t)
    81  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    82  	defer cancel()
    83  
    84  	tests := []struct {
    85  		name    string
    86  		factory func(context.Context) (jsonrpc2.Listener, error)
    87  	}{
    88  		{"tcp", func(ctx context.Context) (jsonrpc2.Listener, error) {
    89  			return jsonrpc2.NetListener(ctx, "tcp", "localhost:0", jsonrpc2.NetListenOptions{})
    90  		}},
    91  		{"pipe", func(ctx context.Context) (jsonrpc2.Listener, error) {
    92  			return jsonrpc2.NetPipe(ctx)
    93  		}},
    94  	}
    95  
    96  	for _, test := range tests {
    97  		t.Run(test.name, func(t *testing.T) {
    98  			fake, err := test.factory(ctx)
    99  			if err != nil {
   100  				t.Fatal(err)
   101  			}
   102  			conn, shutdown, err := newFake(ctx, fake)
   103  			if err != nil {
   104  				t.Fatal(err)
   105  			}
   106  			defer shutdown(ctx)
   107  			var got msg
   108  			if err := conn.Call(ctx, "ping", &msg{"ting"}).Await(ctx, &got); err != nil {
   109  				t.Fatal(err)
   110  			}
   111  			if want := "pong"; got.Msg != want {
   112  				t.Errorf("conn.Call(...): returned %q, want %q", got, want)
   113  			}
   114  		})
   115  	}
   116  }
   117  
   118  func newFake(ctx context.Context, l jsonrpc2.Listener) (*jsonrpc2.Connection, func(context.Context), error) {
   119  	l = jsonrpc2.NewIdleListener(100*time.Millisecond, l)
   120  	server, err := jsonrpc2.Serve(ctx, l, jsonrpc2.ConnectionOptions{
   121  		Handler: fakeHandler{},
   122  	})
   123  	if err != nil {
   124  		return nil, nil, err
   125  	}
   126  
   127  	client, err := jsonrpc2.Dial(ctx,
   128  		l.Dialer(),
   129  		jsonrpc2.ConnectionOptions{
   130  			Handler: fakeHandler{},
   131  		})
   132  	if err != nil {
   133  		return nil, nil, err
   134  	}
   135  	return client, func(ctx context.Context) {
   136  		l.Close()
   137  		client.Close()
   138  		server.Wait()
   139  	}, nil
   140  }