trpc.group/trpc-go/trpc-go@v1.0.3/pool/connpool/pool_test.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package connpool
    15  
    16  import (
    17  	"context"
    18  	"io"
    19  	"testing"
    20  	"time"
    21  
    22  	"github.com/stretchr/testify/assert"
    23  	"trpc.group/trpc-go/trpc-go/codec"
    24  )
    25  
    26  func TestWithGetOptions(t *testing.T) {
    27  	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
    28  	defer cancel()
    29  	fb := &noopFramerBuilder{}
    30  	opts := &GetOptions{CustomReader: codec.NewReader,
    31  		FramerBuilder: fb,
    32  		Ctx:           ctx,
    33  	}
    34  
    35  	localAddr := "127.0.0.1:8080"
    36  	opts.WithLocalAddr(localAddr)
    37  	protocol := "xxx-protocol"
    38  	opts.WithProtocol(protocol)
    39  	opts.WithCustomReader(codec.NewReader)
    40  
    41  	assert.Equal(t, opts.FramerBuilder, fb)
    42  	assert.Equal(t, opts.Ctx, ctx)
    43  	assert.Equal(t, opts.LocalAddr, localAddr)
    44  	assert.Equal(t, protocol, opts.Protocol)
    45  	assert.NotNil(t, opts.CustomReader)
    46  }
    47  
    48  type emptyFramerBuilder struct{}
    49  
    50  // New creates a new Framer.
    51  func (*emptyFramerBuilder) New(io.Reader) codec.Framer {
    52  	return &emptyFramer{}
    53  }
    54  
    55  type emptyFramer struct{}
    56  
    57  // ReadFrame reads the frame.
    58  func (*emptyFramer) ReadFrame() ([]byte, error) {
    59  	return nil, nil
    60  }
    61  
    62  type safeFramerBuilder struct{}
    63  
    64  // New creates a new Framer.
    65  func (*safeFramerBuilder) New(io.Reader) codec.Framer {
    66  	return &safeFramer{}
    67  }
    68  
    69  type safeFramer struct{}
    70  
    71  // ReadFrame reads the frame.
    72  func (*safeFramer) ReadFrame() ([]byte, error) {
    73  	in := []byte("hello world!")
    74  	out := make([]byte, len(in))
    75  	copy(out, in)
    76  	return out, nil
    77  }
    78  
    79  func (*safeFramer) IsSafe() bool {
    80  	return true
    81  }
    82  
    83  func TestGetDialCtx(t *testing.T) {
    84  	opts := &GetOptions{CustomReader: codec.NewReader}
    85  	ctx, cancel := opts.getDialCtx(0)
    86  	assert.NotNil(t, ctx)
    87  	assert.NotNil(t, cancel)
    88  
    89  	ctx, cancel = opts.getDialCtx(defaultDialTimeout)
    90  	assert.NotNil(t, ctx)
    91  	assert.NotNil(t, cancel)
    92  
    93  	backgroundCtx := context.Background()
    94  	opts.WithContext(backgroundCtx)
    95  	ctx, cancel = opts.getDialCtx(defaultDialTimeout)
    96  	assert.NotEqual(t, backgroundCtx, ctx)
    97  	assert.NotNil(t, cancel)
    98  
    99  	timeoutCtx, cancel := context.WithTimeout(context.Background(), time.Second)
   100  	defer cancel()
   101  	opts.WithContext(timeoutCtx)
   102  	ctx, cancel = opts.getDialCtx(defaultDialTimeout)
   103  	assert.Equal(t, timeoutCtx, ctx)
   104  	assert.Nil(t, cancel)
   105  
   106  	opts.WithContext(timeoutCtx)
   107  	opts.WithDialTimeout(200 * time.Millisecond)
   108  	ctx, cancel = opts.getDialCtx(defaultDialTimeout)
   109  	assert.NotEqual(t, timeoutCtx, ctx)
   110  	assert.NotNil(t, cancel)
   111  
   112  	opts.WithContext(timeoutCtx)
   113  	opts.WithDialTimeout(2 * time.Second)
   114  	ctx, cancel = opts.getDialCtx(defaultDialTimeout)
   115  	assert.Equal(t, timeoutCtx, ctx)
   116  	assert.Nil(t, cancel)
   117  }