github.com/lingyao2333/mo-zero@v1.4.1/zrpc/internal/clientinterceptors/timeoutinterceptor_test.go (about)

     1  package clientinterceptors
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"google.golang.org/grpc"
    12  )
    13  
    14  func TestTimeoutInterceptor(t *testing.T) {
    15  	timeouts := []time.Duration{0, time.Millisecond * 10}
    16  	for _, timeout := range timeouts {
    17  		t.Run(strconv.FormatInt(int64(timeout), 10), func(t *testing.T) {
    18  			interceptor := TimeoutInterceptor(timeout)
    19  			cc := new(grpc.ClientConn)
    20  			err := interceptor(context.Background(), "/foo", nil, nil, cc,
    21  				func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
    22  					opts ...grpc.CallOption) error {
    23  					return nil
    24  				},
    25  			)
    26  			assert.Nil(t, err)
    27  		})
    28  	}
    29  }
    30  
    31  func TestTimeoutInterceptor_timeout(t *testing.T) {
    32  	const timeout = time.Millisecond * 10
    33  	interceptor := TimeoutInterceptor(timeout)
    34  	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
    35  	defer cancel()
    36  	var wg sync.WaitGroup
    37  	wg.Add(1)
    38  	cc := new(grpc.ClientConn)
    39  	err := interceptor(ctx, "/foo", nil, nil, cc,
    40  		func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
    41  			opts ...grpc.CallOption) error {
    42  			defer wg.Done()
    43  			tm, ok := ctx.Deadline()
    44  			assert.True(t, ok)
    45  			assert.True(t, tm.Before(time.Now().Add(timeout+time.Millisecond)))
    46  			return nil
    47  		})
    48  	wg.Wait()
    49  	assert.Nil(t, err)
    50  }
    51  
    52  func TestTimeoutInterceptor_panic(t *testing.T) {
    53  	timeouts := []time.Duration{0, time.Millisecond * 10}
    54  	for _, timeout := range timeouts {
    55  		t.Run(strconv.FormatInt(int64(timeout), 10), func(t *testing.T) {
    56  			interceptor := TimeoutInterceptor(timeout)
    57  			cc := new(grpc.ClientConn)
    58  			assert.Panics(t, func() {
    59  				_ = interceptor(context.Background(), "/foo", nil, nil, cc,
    60  					func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
    61  						opts ...grpc.CallOption) error {
    62  						panic("any")
    63  					},
    64  				)
    65  			})
    66  		})
    67  	}
    68  }