...

Source file src/edge-infra.dev/pkg/sds/lib/etcd/client/retry/retry_client_test.go

Documentation: edge-infra.dev/pkg/sds/lib/etcd/client/retry

     1  package retry
     2  
     3  import (
     4  	"context"
     5  	"crypto/tls"
     6  	"fmt"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/go-logr/logr/testr"
    11  	"github.com/golang/mock/gomock"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  	"go.etcd.io/etcd/api/v3/etcdserverpb"
    15  	clientv3 "go.etcd.io/etcd/client/v3"
    16  	ctrl "sigs.k8s.io/controller-runtime"
    17  
    18  	"edge-infra.dev/pkg/sds/lib/etcd/client"
    19  	"edge-infra.dev/pkg/sds/lib/etcd/client/mocks"
    20  )
    21  
    22  var testMember = &etcdserverpb.Member{
    23  	ID:         123,
    24  	Name:       "test-name",
    25  	PeerURLs:   []string{"test-peer-url"},
    26  	ClientURLs: []string{"test-client-url"},
    27  	IsLearner:  true,
    28  }
    29  
    30  var testMemberListResp = &clientv3.MemberListResponse{
    31  	Members: []*etcdserverpb.Member{testMember},
    32  }
    33  
    34  var testMemberAddResp = &clientv3.MemberAddResponse{
    35  	Member: testMember,
    36  }
    37  
    38  var testMemberPromoteResp = &clientv3.MemberPromoteResponse{
    39  	Members: []*etcdserverpb.Member{testMember},
    40  }
    41  
    42  var testMemberRemoveResp = &clientv3.MemberRemoveResponse{
    43  	Members: []*etcdserverpb.Member{testMember},
    44  }
    45  
    46  func SetupTestCtx(t *testing.T) context.Context {
    47  	logOptions := testr.Options{
    48  		LogTimestamp: true,
    49  		Verbosity:    -1,
    50  	}
    51  
    52  	ctx := ctrl.LoggerInto(context.Background(), testr.NewWithOptions(t, logOptions))
    53  	return ctx
    54  }
    55  
    56  func TestNew(t *testing.T) {
    57  	testCases := map[string]struct {
    58  		input    Config
    59  		expected Config
    60  	}{
    61  		"Defaults": {
    62  			input: Config{},
    63  			expected: Config{
    64  				RequestTimeout: 5 * time.Second,
    65  				InitialBackoff: 500 * time.Millisecond,
    66  				BackoffFactor:  1.5,
    67  				MaxRetries:     3,
    68  			},
    69  		},
    70  		"PartialOverwrite": {
    71  			input: Config{
    72  				BackoffFactor: 2,
    73  				MaxRetries:    5,
    74  			},
    75  			expected: Config{
    76  				RequestTimeout: 5 * time.Second,
    77  				InitialBackoff: 500 * time.Millisecond,
    78  				BackoffFactor:  2,
    79  				MaxRetries:     5,
    80  			},
    81  		},
    82  		"Overwrite": {
    83  			input: Config{
    84  				RequestTimeout: 10 * time.Second,
    85  				InitialBackoff: 5 * time.Second,
    86  				BackoffFactor:  2,
    87  				MaxRetries:     5,
    88  			},
    89  			expected: Config{
    90  				RequestTimeout: 10 * time.Second,
    91  				InitialBackoff: 5 * time.Second,
    92  				BackoffFactor:  2,
    93  				MaxRetries:     5,
    94  			},
    95  		},
    96  	}
    97  
    98  	for name, tc := range testCases {
    99  		t.Run(name, func(t *testing.T) {
   100  			rcli := New(*createClient(nil), tc.input)
   101  
   102  			assert.Equal(t, tc.expected, rcli.(*Client).Config)
   103  		})
   104  	}
   105  }
   106  
   107  func TestWithRetry(t *testing.T) {
   108  	var count int
   109  
   110  	testCases := map[string]struct {
   111  		client           Retrier
   112  		fn               func(ctx context.Context) error
   113  		expectError      bool
   114  		expectedAttempts int
   115  	}{
   116  		"Success": {
   117  			client: New(*createClient(nil), Config{}),
   118  			fn: func(_ context.Context) error {
   119  				count++
   120  				return nil
   121  			},
   122  			expectError:      false,
   123  			expectedAttempts: 1,
   124  		},
   125  		"SuccessAfterRetry": {
   126  			client: New(*createClient(nil), Config{}),
   127  			fn: func(_ context.Context) error {
   128  				count++
   129  				if count < 3 {
   130  					return fmt.Errorf("force retry")
   131  				}
   132  				return nil
   133  			},
   134  			expectError:      false,
   135  			expectedAttempts: 3,
   136  		},
   137  		"FailureAfterRetry": {
   138  			client: New(*createClient(nil), Config{
   139  				MaxRetries: 2,
   140  			}),
   141  			fn: func(_ context.Context) error {
   142  				count++
   143  				return fmt.Errorf("example failure")
   144  			},
   145  			expectError:      true,
   146  			expectedAttempts: 3,
   147  		},
   148  	}
   149  
   150  	for name, tc := range testCases {
   151  		t.Run(name, func(t *testing.T) {
   152  			count = 0
   153  			cli := tc.client.(*Client)
   154  
   155  			err := cli.withRetry(SetupTestCtx(t), tc.fn)
   156  			switch tc.expectError {
   157  			case true:
   158  				assert.Error(t, err)
   159  			case false:
   160  				assert.NoError(t, err)
   161  			}
   162  
   163  			assert.Equal(t, tc.expectedAttempts, count)
   164  		})
   165  	}
   166  }
   167  
   168  func TestSafeMemberList(t *testing.T) {
   169  	mockCtrl := gomock.NewController(t)
   170  	defer mockCtrl.Finish()
   171  
   172  	retryClient := getSafeMemberListMockClient(mockCtrl)
   173  	resp, err := retryClient.SafeMemberList(SetupTestCtx(t))
   174  	require.NoError(t, err)
   175  	require.Equal(t, resp, testMemberListResp)
   176  }
   177  
   178  func TestSafeMemberAddAsLearner(t *testing.T) {
   179  	mockCtrl := gomock.NewController(t)
   180  	defer mockCtrl.Finish()
   181  
   182  	retryClient := getSafeMemberAddAsLearnerMockClient(mockCtrl)
   183  	resp, err := retryClient.SafeMemberAddAsLearner(SetupTestCtx(t), []string{})
   184  	require.NoError(t, err)
   185  	require.Equal(t, resp, testMemberAddResp)
   186  }
   187  
   188  func TestSafeMemberPromote(t *testing.T) {
   189  	mockCtrl := gomock.NewController(t)
   190  	defer mockCtrl.Finish()
   191  
   192  	retryClient := getSafeMemberPromoteMockClient(mockCtrl)
   193  	resp, err := retryClient.SafeMemberPromote(SetupTestCtx(t), 0)
   194  	require.NoError(t, err)
   195  	require.Equal(t, resp, testMemberPromoteResp)
   196  }
   197  
   198  func TestSafeMemberRemove(t *testing.T) {
   199  	mockCtrl := gomock.NewController(t)
   200  	defer mockCtrl.Finish()
   201  
   202  	retryClient := getSafeMemberRemoveMockClient(mockCtrl)
   203  	resp, err := retryClient.SafeMemberRemove(SetupTestCtx(t), 0)
   204  	require.NoError(t, err)
   205  	require.Equal(t, resp, testMemberRemoveResp)
   206  }
   207  
   208  func getSafeMemberListMockClient(mockCtrl *gomock.Controller) Retrier {
   209  	mockCluster := mocks.NewMockCluster(mockCtrl)
   210  	mockCluster.EXPECT().MemberList(gomock.Any()).DoAndReturn(func(_ context.Context) (resp *clientv3.MemberListResponse, err error) {
   211  		resp = testMemberListResp
   212  		return resp, nil
   213  	})
   214  
   215  	return New(*createClient(mockCluster), Config{})
   216  }
   217  
   218  func getSafeMemberAddAsLearnerMockClient(mockCtrl *gomock.Controller) Retrier {
   219  	mockCluster := mocks.NewMockCluster(mockCtrl)
   220  	mockCluster.EXPECT().MemberAddAsLearner(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ []string) (resp *clientv3.MemberAddResponse, err error) {
   221  		resp = testMemberAddResp
   222  		return resp, nil
   223  	})
   224  
   225  	return New(*createClient(mockCluster), Config{})
   226  }
   227  
   228  func getSafeMemberPromoteMockClient(mockCtrl *gomock.Controller) Retrier {
   229  	mockCluster := mocks.NewMockCluster(mockCtrl)
   230  	mockCluster.EXPECT().MemberPromote(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ uint64) (resp *clientv3.MemberPromoteResponse, err error) {
   231  		resp = testMemberPromoteResp
   232  		return resp, nil
   233  	})
   234  
   235  	return New(*createClient(mockCluster), Config{})
   236  }
   237  
   238  func getSafeMemberRemoveMockClient(mockCtrl *gomock.Controller) Retrier {
   239  	mockCluster := mocks.NewMockCluster(mockCtrl)
   240  	mockCluster.EXPECT().MemberRemove(gomock.Any(), gomock.Any()).DoAndReturn(func(_ context.Context, _ uint64) (resp *clientv3.MemberRemoveResponse, err error) {
   241  		resp = testMemberRemoveResp
   242  		return resp, nil
   243  	})
   244  
   245  	return New(*createClient(mockCluster), Config{})
   246  }
   247  
   248  func createClient(mockCluster *mocks.MockCluster) *clientv3.Client {
   249  	config := &tls.Config{
   250  		MinVersion: tls.VersionTLS12,
   251  	}
   252  	cli, err := client.New(config, 5*time.Second, "fake-ip:0")
   253  	if err != nil {
   254  		fmt.Println(err)
   255  	}
   256  	cli.Cluster = mockCluster
   257  
   258  	return cli
   259  }
   260  

View as plain text