github.com/grafana/pyroscope@v1.18.0/pkg/util/delayhandler/connect_test.go (about)

     1  package delayhandler
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"connectrpc.com/connect"
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/grafana/pyroscope/pkg/tenant"
    13  )
    14  
    15  func TestConnectInterceptor(t *testing.T) {
    16  	now := time.Unix(1718211600, 0)
    17  	tenantID := "tenant"
    18  
    19  	tests := []struct {
    20  		name         string
    21  		configDelay  time.Duration
    22  		cancelDelay  bool
    23  		expectSleep  bool
    24  		expectHeader bool
    25  	}{
    26  		{
    27  			name:         "no delay",
    28  			configDelay:  0,
    29  			expectSleep:  false,
    30  			expectHeader: false,
    31  		},
    32  		{
    33  			name:         "with delay",
    34  			configDelay:  100 * time.Millisecond,
    35  			expectSleep:  true,
    36  			expectHeader: true,
    37  		},
    38  		{
    39  			name:         "cancelled delay",
    40  			configDelay:  100 * time.Millisecond,
    41  			cancelDelay:  true,
    42  			expectSleep:  false,
    43  			expectHeader: false,
    44  		},
    45  	}
    46  
    47  	for _, tt := range tests {
    48  		t.Run(tt.name, func(t *testing.T) {
    49  			timeNowMock(t, []time.Time{now, now.Add(5 * time.Millisecond)})
    50  			sleeps, cleanUpSleep := timeAfterMock()
    51  			defer cleanUpSleep()
    52  
    53  			limits := newMockLimits()
    54  			limits.setDelay(tenantID, tt.configDelay)
    55  
    56  			interceptor := NewConnect(limits)
    57  
    58  			handler := connect.UnaryFunc(func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
    59  				if tt.cancelDelay {
    60  					CancelDelay(ctx)
    61  				}
    62  				return connect.NewResponse(&struct{}{}), nil
    63  			})
    64  
    65  			wrappedHandler := interceptor.WrapUnary(handler)
    66  			req := connect.NewRequest(&struct{}{})
    67  			ctx := tenant.InjectTenantID(context.Background(), tenantID)
    68  
    69  			resp, err := wrappedHandler(ctx, req)
    70  
    71  			require.NoError(t, err)
    72  			require.NotNil(t, resp)
    73  
    74  			if tt.expectSleep {
    75  				require.Len(t, sleeps.values, 1)
    76  				assert.Greater(t, sleeps.values[0], 80*time.Millisecond)
    77  				assert.Less(t, sleeps.values[0], 120*time.Millisecond)
    78  			} else {
    79  				require.Len(t, sleeps.values, 0)
    80  			}
    81  
    82  			if tt.expectHeader {
    83  				assert.Contains(t, resp.Header().Get("Server-Timing"), "artificial_delay")
    84  			} else {
    85  				assert.Empty(t, resp.Header().Get("Server-Timing"))
    86  			}
    87  		})
    88  	}
    89  }