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

     1  package serverinterceptors
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"google.golang.org/grpc"
    11  	"google.golang.org/grpc/codes"
    12  	"google.golang.org/grpc/status"
    13  )
    14  
    15  func TestUnaryTimeoutInterceptor(t *testing.T) {
    16  	interceptor := UnaryTimeoutInterceptor(time.Millisecond * 10)
    17  	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
    18  		FullMethod: "/",
    19  	}, func(ctx context.Context, req interface{}) (interface{}, error) {
    20  		return nil, nil
    21  	})
    22  	assert.Nil(t, err)
    23  }
    24  
    25  func TestUnaryTimeoutInterceptor_panic(t *testing.T) {
    26  	interceptor := UnaryTimeoutInterceptor(time.Millisecond * 10)
    27  	assert.Panics(t, func() {
    28  		_, _ = interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
    29  			FullMethod: "/",
    30  		}, func(ctx context.Context, req interface{}) (interface{}, error) {
    31  			panic("any")
    32  		})
    33  	})
    34  }
    35  
    36  func TestUnaryTimeoutInterceptor_timeout(t *testing.T) {
    37  	const timeout = time.Millisecond * 10
    38  	interceptor := UnaryTimeoutInterceptor(timeout)
    39  	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
    40  	defer cancel()
    41  	var wg sync.WaitGroup
    42  	wg.Add(1)
    43  	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
    44  		FullMethod: "/",
    45  	}, func(ctx context.Context, req interface{}) (interface{}, error) {
    46  		defer wg.Done()
    47  		tm, ok := ctx.Deadline()
    48  		assert.True(t, ok)
    49  		assert.True(t, tm.Before(time.Now().Add(timeout+time.Millisecond)))
    50  		return nil, nil
    51  	})
    52  	wg.Wait()
    53  	assert.Nil(t, err)
    54  }
    55  
    56  func TestUnaryTimeoutInterceptor_timeoutExpire(t *testing.T) {
    57  	const timeout = time.Millisecond * 10
    58  	interceptor := UnaryTimeoutInterceptor(timeout)
    59  	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
    60  	defer cancel()
    61  	var wg sync.WaitGroup
    62  	wg.Add(1)
    63  	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
    64  		FullMethod: "/",
    65  	}, func(ctx context.Context, req interface{}) (interface{}, error) {
    66  		defer wg.Done()
    67  		time.Sleep(time.Millisecond * 50)
    68  		return nil, nil
    69  	})
    70  	wg.Wait()
    71  	assert.EqualValues(t, status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error()), err)
    72  }
    73  
    74  func TestUnaryTimeoutInterceptor_cancel(t *testing.T) {
    75  	const timeout = time.Minute * 10
    76  	interceptor := UnaryTimeoutInterceptor(timeout)
    77  	ctx, cancel := context.WithCancel(context.Background())
    78  	cancel()
    79  
    80  	var wg sync.WaitGroup
    81  	wg.Add(1)
    82  	_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
    83  		FullMethod: "/",
    84  	}, func(ctx context.Context, req interface{}) (interface{}, error) {
    85  		defer wg.Done()
    86  		time.Sleep(time.Millisecond * 50)
    87  		return nil, nil
    88  	})
    89  
    90  	wg.Wait()
    91  	assert.EqualValues(t, status.Error(codes.Canceled, context.Canceled.Error()), err)
    92  }