github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/chunk/client/gcp/gcs_object_client_test.go (about)

     1  package gcp
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  	"time"
    10  
    11  	"cloud.google.com/go/storage"
    12  	"github.com/stretchr/testify/require"
    13  	"go.uber.org/atomic"
    14  	"google.golang.org/api/option"
    15  
    16  	"github.com/grafana/loki/pkg/storage/chunk/client/hedging"
    17  )
    18  
    19  func Test_Hedging(t *testing.T) {
    20  	for _, tc := range []struct {
    21  		name          string
    22  		expectedCalls int32
    23  		hedgeAt       time.Duration
    24  		upTo          int
    25  		do            func(c *GCSObjectClient)
    26  	}{
    27  		{
    28  			"delete/put/list are not hedged",
    29  			3,
    30  			20 * time.Nanosecond,
    31  			10,
    32  			func(c *GCSObjectClient) {
    33  				_ = c.DeleteObject(context.Background(), "foo")
    34  				_, _, _ = c.List(context.Background(), "foo", "/")
    35  				_ = c.PutObject(context.Background(), "foo", bytes.NewReader([]byte("bar")))
    36  			},
    37  		},
    38  		{
    39  			"gets are hedged",
    40  			3,
    41  			20 * time.Nanosecond,
    42  			3,
    43  			func(c *GCSObjectClient) {
    44  				_, _, _ = c.GetObject(context.Background(), "foo")
    45  			},
    46  		},
    47  		{
    48  			"gets are not hedged when not configured",
    49  			1,
    50  			0,
    51  			0,
    52  			func(c *GCSObjectClient) {
    53  				_, _, _ = c.GetObject(context.Background(), "foo")
    54  			},
    55  		},
    56  	} {
    57  		tc := tc
    58  		t.Run(tc.name, func(t *testing.T) {
    59  			count := atomic.NewInt32(0)
    60  			server := fakeServer(t, 200*time.Millisecond, count)
    61  			ctx := context.Background()
    62  			c, err := newGCSObjectClient(ctx, GCSConfig{
    63  				BucketName: "test-bucket",
    64  				Insecure:   true,
    65  			}, hedging.Config{
    66  				At:           tc.hedgeAt,
    67  				UpTo:         tc.upTo,
    68  				MaxPerSecond: 1000,
    69  			}, func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) {
    70  				opts = append(opts, option.WithEndpoint(server.URL))
    71  				opts = append(opts, option.WithoutAuthentication())
    72  				return storage.NewClient(ctx, opts...)
    73  			})
    74  			require.NoError(t, err)
    75  			tc.do(c)
    76  			require.Equal(t, tc.expectedCalls, count.Load())
    77  		})
    78  	}
    79  }
    80  
    81  func fakeServer(t *testing.T, returnIn time.Duration, counter *atomic.Int32) *httptest.Server {
    82  	server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    83  		counter.Inc()
    84  		time.Sleep(returnIn)
    85  		_, _ = w.Write([]byte(`{}`))
    86  	}))
    87  	server.StartTLS()
    88  	t.Cleanup(server.Close)
    89  
    90  	return server
    91  }