1 package retryclient
2
3 import (
4 "context"
5 "fmt"
6 "testing"
7
8 "github.com/go-logr/logr/testr"
9 "github.com/golang/mock/gomock"
10 "github.com/stretchr/testify/assert"
11 "github.com/stretchr/testify/require"
12 corev1 "k8s.io/api/core/v1"
13 apierrors "k8s.io/apimachinery/pkg/api/errors"
14 metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
15 kruntime "k8s.io/apimachinery/pkg/runtime"
16 k8stypes "k8s.io/apimachinery/pkg/types"
17 utilruntime "k8s.io/apimachinery/pkg/util/runtime"
18 clientgoscheme "k8s.io/client-go/kubernetes/scheme"
19 ctrl "sigs.k8s.io/controller-runtime"
20 "sigs.k8s.io/controller-runtime/pkg/client"
21 "sigs.k8s.io/controller-runtime/pkg/client/fake"
22
23 "edge-infra.dev/pkg/sds/lib/k8s/retryclient/types"
24 )
25
26 var keyObj = corev1.Pod{
27 ObjectMeta: metav1.ObjectMeta{
28 Name: "test-pod",
29 Namespace: "test-namespace",
30 },
31 }
32
33 func testPodList(pods ...corev1.Pod) *corev1.PodList {
34 return &corev1.PodList{
35 Items: pods,
36 }
37 }
38
39 func testPod(variant string) *corev1.Pod {
40 return &corev1.Pod{TypeMeta: metav1.TypeMeta{
41 Kind: "Pod",
42 APIVersion: "v1",
43 },
44 ObjectMeta: metav1.ObjectMeta{
45 Name: "test-pod",
46 Namespace: "test-namespace",
47 },
48 Spec: corev1.PodSpec{
49 Containers: []corev1.Container{{
50 Name: variant,
51 Image: variant,
52 }},
53 },
54 }
55 }
56
57 func SetupTestCtx(t *testing.T) context.Context {
58 logOptions := testr.Options{
59 LogTimestamp: true,
60 Verbosity: -1,
61 }
62
63 ctx := ctrl.LoggerInto(context.Background(), testr.NewWithOptions(t, logOptions))
64 return ctx
65 }
66
67 func TestWithRetry(t *testing.T) {
68 var count int
69
70 testCases := map[string]struct {
71 client types.Retrier
72 fn func(ctx context.Context) error
73 expectError bool
74 expectedAttempts int
75 }{
76 "Success": {
77 client: New(getMockKubeClient(), getMockKubeReader(), Config{}),
78 fn: func(_ context.Context) error {
79 count++
80 return nil
81 },
82 expectError: false,
83 expectedAttempts: 1,
84 },
85 "SuccessAfterRetry": {
86 client: New(getMockKubeClient(), getMockKubeReader(), Config{}),
87 fn: func(_ context.Context) error {
88 count++
89 if count < 3 {
90 return fmt.Errorf("force retry")
91 }
92 return nil
93 },
94 expectError: false,
95 expectedAttempts: 3,
96 },
97 "FailureAfterRetry": {
98 client: New(getMockKubeClient(), getMockKubeReader(), Config{
99 MaxRetries: 2,
100 }),
101 fn: func(_ context.Context) error {
102 count++
103 return fmt.Errorf("example failure")
104 },
105 expectError: true,
106 expectedAttempts: 3,
107 },
108 }
109
110 for name, tc := range testCases {
111 t.Run(name, func(t *testing.T) {
112 count = 0
113 cli := tc.client.(*RetryClient)
114
115 err := cli.withRetry(SetupTestCtx(t), tc.fn)
116 switch tc.expectError {
117 case true:
118 assert.Error(t, err)
119 case false:
120 assert.NoError(t, err)
121 }
122
123 assert.Equal(t, tc.expectedAttempts, count)
124 })
125 }
126 }
127
128 func TestSafeGet(t *testing.T) {
129 mockCtrl := gomock.NewController(t)
130 defer mockCtrl.Finish()
131
132 cli := New(
133 getMockKubeClient(testPod("cached")),
134 getMockKubeReader(testPod("uncached")),
135 Config{})
136
137 testCases := map[string]struct {
138 fn func(context.Context, client.Object) error
139 expected *corev1.Pod
140 }{
141 "UseCache": {
142 fn: func(ctx context.Context, obj client.Object) error {
143 return cli.SafeGet(ctx, client.ObjectKeyFromObject(obj), obj)
144 },
145 expected: testPod("cached"),
146 },
147 "IgnoreCache": {
148 fn: func(ctx context.Context, obj client.Object) error {
149 return cli.IgnoreCache().SafeGet(ctx, client.ObjectKeyFromObject(obj), obj)
150 },
151 expected: testPod("uncached"),
152 },
153 }
154
155 for name, tc := range testCases {
156 t.Run(name, func(t *testing.T) {
157 obj := keyObj.DeepCopy()
158 require.NoError(t, tc.fn(SetupTestCtx(t), obj))
159 obj.ResourceVersion = ""
160 assert.Equal(t, tc.expected, obj)
161 })
162 }
163 }
164
165 func TestSafeList(t *testing.T) {
166 mockCtrl := gomock.NewController(t)
167 defer mockCtrl.Finish()
168
169 cli := New(
170 getMockKubeClient(testPod("cached")),
171 getMockKubeReader(testPod("uncached")),
172 Config{})
173
174 testCases := map[string]struct {
175 fn func(context.Context, *corev1.PodList) error
176 expected *corev1.PodList
177 }{
178 "UseCache": {
179 fn: func(ctx context.Context, podList *corev1.PodList) error {
180 return cli.SafeList(ctx, podList)
181 },
182 expected: testPodList(*testPod("cached")),
183 },
184 "IgnoreCache": {
185 fn: func(ctx context.Context, podList *corev1.PodList) error {
186 return cli.IgnoreCache().SafeList(ctx, podList)
187 },
188 expected: testPodList(*testPod("uncached")),
189 },
190 }
191
192 for name, tc := range testCases {
193 t.Run(name, func(t *testing.T) {
194 podList := &corev1.PodList{
195 TypeMeta: metav1.TypeMeta{
196 Kind: "PodList",
197 APIVersion: "v1",
198 },
199 }
200 require.NoError(t, tc.fn(SetupTestCtx(t), podList))
201 for i := range podList.Items {
202 podList.Items[i].ResourceVersion = ""
203 }
204 assert.Equal(t, tc.expected, podList)
205 })
206 }
207 }
208
209 func TestSafeCreate(t *testing.T) {
210 mockCtrl := gomock.NewController(t)
211 defer mockCtrl.Finish()
212
213 cli := New(getMockKubeClient(), getMockKubeReader(), Config{})
214
215 require.NoError(t, cli.SafeCreate(SetupTestCtx(t), testPod("test-image")))
216 p := keyObj.DeepCopy()
217 err := cli.Client().Get(SetupTestCtx(t), client.ObjectKeyFromObject(keyObj.DeepCopy()), p)
218 require.NoError(t, err)
219 p.ResourceVersion = ""
220 assert.Equal(t, testPod("test-image"), p)
221 }
222
223 func TestSafeDelete(t *testing.T) {
224 mockCtrl := gomock.NewController(t)
225 defer mockCtrl.Finish()
226
227 cli := New(
228 getMockKubeClient(testPod("test-image")),
229 getMockKubeReader(),
230 Config{})
231
232 require.NoError(t, cli.SafeDelete(SetupTestCtx(t), testPod("test-image")))
233 err := cli.Client().Get(SetupTestCtx(t), client.ObjectKeyFromObject(keyObj.DeepCopy()), keyObj.DeepCopy())
234 assert.True(t, apierrors.IsNotFound(err))
235 }
236
237 func TestSafeUpdate(t *testing.T) {
238 mockCtrl := gomock.NewController(t)
239 defer mockCtrl.Finish()
240
241 cli := New(
242 getMockKubeClient(testPod("test-image-1")),
243 getMockKubeReader(testPod("test-image-2")),
244 Config{})
245
246 testCases := map[string]struct {
247 fn func(context.Context, k8stypes.NamespacedName, client.Object, func(context.Context, client.Object) error, ...client.UpdateOption) error
248 updateFn func(context.Context, client.Object) error
249 expected *corev1.Pod
250 }{
251 "UseCache": {
252 fn: cli.SafeUpdate,
253 updateFn: func(_ context.Context, obj client.Object) error {
254 testPod("test-image-2").DeepCopyInto(obj.(*corev1.Pod))
255 return nil
256 },
257 expected: testPod("test-image-2"),
258 },
259 "IgnoreCache": {
260 fn: cli.IgnoreCache().SafeUpdate,
261 updateFn: func(_ context.Context, obj client.Object) error {
262 testPod("test-image-1").DeepCopyInto(obj.(*corev1.Pod))
263 return nil
264 },
265 expected: testPod("test-image-1"),
266 },
267 }
268
269 for name, tc := range testCases {
270 t.Run(name, func(t *testing.T) {
271 pod := &corev1.Pod{}
272 err := tc.fn(SetupTestCtx(t), client.ObjectKeyFromObject(keyObj.DeepCopy()), pod, tc.updateFn)
273 require.NoError(t, err)
274 pod.ResourceVersion = ""
275 assert.Equal(t, tc.expected, pod)
276 })
277 }
278 }
279
280 func TestIgnoreCache(t *testing.T) {
281 mockCtrl := gomock.NewController(t)
282 defer mockCtrl.Finish()
283
284 cli := New(
285 getMockKubeClient(),
286 getMockKubeReader(),
287 Config{})
288 rcli := cli.(*RetryClient)
289
290
291 assert.True(t, rcli.useCache)
292 assert.False(t, rcli.IgnoreCache().(*RetryClient).useCache)
293 assert.True(t, rcli.useCache)
294 }
295
296
297 func getMockKubeClient(initObjs ...client.Object) client.Client {
298 return fake.NewClientBuilder().WithScheme(createScheme()).WithObjects(initObjs...).Build()
299 }
300
301 func getMockKubeReader(initObjs ...client.Object) client.Reader {
302 return fake.NewClientBuilder().WithScheme(createScheme()).WithObjects(initObjs...).Build()
303 }
304
305
306 func createScheme() *kruntime.Scheme {
307 scheme := kruntime.NewScheme()
308 utilruntime.Must(clientgoscheme.AddToScheme(scheme))
309 return scheme
310 }
311
View as plain text