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

     1  package openstack
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"net/http"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  	"go.uber.org/atomic"
    12  
    13  	"github.com/grafana/loki/pkg/storage/bucket/swift"
    14  	"github.com/grafana/loki/pkg/storage/chunk/client/hedging"
    15  )
    16  
    17  type RoundTripperFunc func(*http.Request) (*http.Response, error)
    18  
    19  func (fn RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
    20  	return fn(req)
    21  }
    22  
    23  func Test_Hedging(t *testing.T) {
    24  	for _, tc := range []struct {
    25  		name          string
    26  		expectedCalls int32
    27  		hedgeAt       time.Duration
    28  		upTo          int
    29  		do            func(c *SwiftObjectClient)
    30  	}{
    31  		{
    32  			"delete/put/list are not hedged",
    33  			3,
    34  			20 * time.Nanosecond,
    35  			10,
    36  			func(c *SwiftObjectClient) {
    37  				_ = c.DeleteObject(context.Background(), "foo")
    38  				_, _, _ = c.List(context.Background(), "foo", "/")
    39  				_ = c.PutObject(context.Background(), "foo", bytes.NewReader([]byte("bar")))
    40  			},
    41  		},
    42  		{
    43  			"gets are hedged",
    44  			3,
    45  			20 * time.Nanosecond,
    46  			3,
    47  			func(c *SwiftObjectClient) {
    48  				_, _, _ = c.GetObject(context.Background(), "foo")
    49  			},
    50  		},
    51  		{
    52  			"gets are not hedged when not configured",
    53  			1,
    54  			0,
    55  			0,
    56  			func(c *SwiftObjectClient) {
    57  				_, _, _ = c.GetObject(context.Background(), "foo")
    58  			},
    59  		},
    60  	} {
    61  		tc := tc
    62  		t.Run(tc.name, func(t *testing.T) {
    63  			count := atomic.NewInt32(0)
    64  			// hijack the transport to count the number of calls
    65  			defaultTransport = RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
    66  				// fake auth
    67  				if req.Header.Get("X-Auth-Key") == "passwd" {
    68  					return &http.Response{
    69  						StatusCode: http.StatusOK,
    70  						Body:       http.NoBody,
    71  						Header: http.Header{
    72  							"X-Storage-Url": []string{"http://swift.example.com/v1/AUTH_test"},
    73  							"X-Auth-Token":  []string{"token"},
    74  						},
    75  					}, nil
    76  				}
    77  				// fake container creation
    78  				if req.Method == "PUT" && req.URL.Path == "/v1/AUTH_test/foo" {
    79  					return &http.Response{
    80  						StatusCode: http.StatusCreated,
    81  						Body:       http.NoBody,
    82  					}, nil
    83  				}
    84  				count.Inc()
    85  				time.Sleep(200 * time.Millisecond)
    86  				return &http.Response{
    87  					StatusCode: http.StatusOK,
    88  					Body:       http.NoBody,
    89  				}, nil
    90  			})
    91  
    92  			c, err := NewSwiftObjectClient(SwiftConfig{
    93  				Config: swift.Config{
    94  					MaxRetries:     1,
    95  					ContainerName:  "foo",
    96  					AuthVersion:    1,
    97  					Password:       "passwd",
    98  					ConnectTimeout: 10 * time.Second,
    99  					RequestTimeout: 10 * time.Second,
   100  				},
   101  			}, hedging.Config{
   102  				At:           tc.hedgeAt,
   103  				UpTo:         tc.upTo,
   104  				MaxPerSecond: 1000,
   105  			})
   106  			require.NoError(t, err)
   107  			tc.do(c)
   108  			require.Equal(t, tc.expectedCalls, count.Load())
   109  		})
   110  	}
   111  }