@@ -140,7 +140,8 @@ const (
140
140
// TestCluster wraps clusters with their individual ClientSets so that helper methods can be called.
141
141
type TestCluster struct {
142
142
clusterName string
143
- client kubernetes.Interface
143
+
144
+ client KubernetesClient
144
145
145
146
// testNodepoolRuntimeOverride, if set, overrides the runtime used for pods
146
147
// running on the test nodepool. If unset, the test nodepool's default
@@ -209,6 +210,12 @@ func NewTestClusterFromProto(ctx context.Context, cluster *testpb.Cluster) (*Tes
209
210
210
211
// NewTestClusterFromClient returns a new TestCluster client with a given client.
211
212
func NewTestClusterFromClient (clusterName string , client kubernetes.Interface ) * TestCluster {
213
+ return NewTestClusterFromKubernetesClient (clusterName , & simpleClient {client })
214
+ }
215
+
216
+ // NewTestClusterFromKubernetesClient returns a new TestCluster client with a
217
+ // given KubernetesClient.
218
+ func NewTestClusterFromKubernetesClient (clusterName string , client KubernetesClient ) * TestCluster {
212
219
return & TestCluster {
213
220
clusterName : clusterName ,
214
221
client : client ,
@@ -248,17 +255,24 @@ func (t *TestCluster) OverrideTestNodepoolRuntime(testRuntime RuntimeType) {
248
255
249
256
// createNamespace creates a namespace.
250
257
func (t * TestCluster ) createNamespace (ctx context.Context , namespace * v13.Namespace ) (* v13.Namespace , error ) {
251
- return t .client .CoreV1 ().Namespaces ().Create (ctx , namespace , v1.CreateOptions {})
258
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Namespace , error ) {
259
+ return client .CoreV1 ().Namespaces ().Create (ctx , namespace , v1.CreateOptions {})
260
+ })
252
261
}
253
262
254
263
// getNamespace returns the given namespace in the cluster if it exists.
255
264
func (t * TestCluster ) getNamespace (ctx context.Context , namespaceName string ) (* v13.Namespace , error ) {
256
- return t .client .CoreV1 ().Namespaces ().Get (ctx , namespaceName , v1.GetOptions {})
265
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Namespace , error ) {
266
+ return client .CoreV1 ().Namespaces ().Get (ctx , namespaceName , v1.GetOptions {})
267
+ })
257
268
}
258
269
259
270
// deleteNamespace is a helper method to delete a namespace.
260
271
func (t * TestCluster ) deleteNamespace (ctx context.Context , namespaceName string ) error {
261
- if err := t .client .CoreV1 ().Namespaces ().Delete (ctx , namespaceName , v1.DeleteOptions {}); err != nil {
272
+ err := t .client .Do (ctx , func (ctx context.Context , client kubernetes.Interface ) error {
273
+ return client .CoreV1 ().Namespaces ().Delete (ctx , namespaceName , v1.DeleteOptions {})
274
+ })
275
+ if err != nil {
262
276
return err
263
277
}
264
278
// Wait for the namespace to disappear or for the context to expire.
@@ -282,7 +296,9 @@ func (t *TestCluster) getNodePool(ctx context.Context, nodepoolType NodePoolType
282
296
t .nodepoolsMu .Lock ()
283
297
defer t .nodepoolsMu .Unlock ()
284
298
if t .nodepools == nil {
285
- nodes , err := t .client .CoreV1 ().Nodes ().List (ctx , v1.ListOptions {})
299
+ nodes , err := request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.NodeList , error ) {
300
+ return client .CoreV1 ().Nodes ().List (ctx , v1.ListOptions {})
301
+ })
286
302
if err != nil {
287
303
return nil , fmt .Errorf ("cannot list nodes: %w" , err )
288
304
}
@@ -363,30 +379,39 @@ func (t *TestCluster) HasGVisorTestRuntime(ctx context.Context) (bool, error) {
363
379
if err != nil {
364
380
return false , err
365
381
}
366
- return testNodePool .runtime == RuntimeTypeGVisor || testNodePool .runtime == RuntimeTypeGVisorNvidia , nil
382
+ return testNodePool .runtime == RuntimeTypeGVisor || testNodePool .runtime == RuntimeTypeGVisorTPU , nil
367
383
}
368
384
369
385
// CreatePod is a helper to create a pod.
370
386
func (t * TestCluster ) CreatePod (ctx context.Context , pod * v13.Pod ) (* v13.Pod , error ) {
371
387
if pod .GetObjectMeta ().GetNamespace () == "" {
372
388
pod .SetNamespace (NamespaceDefault )
373
389
}
374
- return t .client .CoreV1 ().Pods (pod .GetNamespace ()).Create (ctx , pod , v1.CreateOptions {})
390
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Pod , error ) {
391
+ return client .CoreV1 ().Pods (pod .GetNamespace ()).Create (ctx , pod , v1.CreateOptions {})
392
+ })
375
393
}
376
394
377
395
// GetPod is a helper method to Get a pod's metadata.
378
396
func (t * TestCluster ) GetPod (ctx context.Context , pod * v13.Pod ) (* v13.Pod , error ) {
379
- return t .client .CoreV1 ().Pods (pod .GetNamespace ()).Get (ctx , pod .GetName (), v1.GetOptions {})
397
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Pod , error ) {
398
+ return client .CoreV1 ().Pods (pod .GetNamespace ()).Get (ctx , pod .GetName (), v1.GetOptions {})
399
+ })
380
400
}
381
401
382
402
// ListPods is a helper method to List pods in a cluster.
383
403
func (t * TestCluster ) ListPods (ctx context.Context , namespace string ) (* v13.PodList , error ) {
384
- return t .client .CoreV1 ().Pods (namespace ).List (ctx , v1.ListOptions {})
404
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.PodList , error ) {
405
+ return client .CoreV1 ().Pods (namespace ).List (ctx , v1.ListOptions {})
406
+ })
385
407
}
386
408
387
409
// DeletePod is a helper method to delete a pod.
388
410
func (t * TestCluster ) DeletePod (ctx context.Context , pod * v13.Pod ) error {
389
- if err := t .client .CoreV1 ().Pods (pod .GetNamespace ()).Delete (ctx , pod .GetName (), v1.DeleteOptions {}); err != nil {
411
+ err := t .client .Do (ctx , func (ctx context.Context , client kubernetes.Interface ) error {
412
+ return client .CoreV1 ().Pods (pod .GetNamespace ()).Delete (ctx , pod .GetName (), v1.DeleteOptions {})
413
+ })
414
+ if err != nil {
390
415
return err
391
416
}
392
417
// Wait for the pod to disappear or for the context to expire.
@@ -406,7 +431,9 @@ func (t *TestCluster) DeletePod(ctx context.Context, pod *v13.Pod) error {
406
431
// GetLogReader gets an io.ReadCloser from which logs can be read. It is the caller's
407
432
// responsibility to close it.
408
433
func (t * TestCluster ) GetLogReader (ctx context.Context , pod * v13.Pod , opts v13.PodLogOptions ) (io.ReadCloser , error ) {
409
- return t .client .CoreV1 ().Pods (pod .GetNamespace ()).GetLogs (pod .GetName (), & opts ).Stream (ctx )
434
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (io.ReadCloser , error ) {
435
+ return client .CoreV1 ().Pods (pod .GetNamespace ()).GetLogs (pod .GetName (), & opts ).Stream (ctx )
436
+ })
410
437
}
411
438
412
439
// ReadPodLogs reads logs from a pod.
@@ -602,22 +629,36 @@ func (t *TestCluster) ContainerDurationSecondsByName(ctx context.Context, pod *v
602
629
603
630
// CreateService is a helper method to create a service in a cluster.
604
631
func (t * TestCluster ) CreateService (ctx context.Context , service * v13.Service ) (* v13.Service , error ) {
605
- return t .client .CoreV1 ().Services (service .GetNamespace ()).Create (ctx , service , v1.CreateOptions {})
632
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Service , error ) {
633
+ return client .CoreV1 ().Services (service .GetNamespace ()).Create (ctx , service , v1.CreateOptions {})
634
+ })
635
+ }
636
+
637
+ // GetService is a helper method to get a service in a cluster.
638
+ func (t * TestCluster ) GetService (ctx context.Context , service * v13.Service ) (* v13.Service , error ) {
639
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.Service , error ) {
640
+ return client .CoreV1 ().Services (service .GetNamespace ()).Get (ctx , service .GetName (), v1.GetOptions {})
641
+ })
606
642
}
607
643
608
644
// ListServices is a helper method to List services in a cluster.
609
645
func (t * TestCluster ) ListServices (ctx context.Context , namespace string ) (* v13.ServiceList , error ) {
610
- return t .client .CoreV1 ().Services (namespace ).List (ctx , v1.ListOptions {})
646
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.ServiceList , error ) {
647
+ return client .CoreV1 ().Services (namespace ).List (ctx , v1.ListOptions {})
648
+ })
611
649
}
612
650
613
651
// DeleteService is a helper to delete a given service.
614
652
func (t * TestCluster ) DeleteService (ctx context.Context , service * v13.Service ) error {
615
- if err := t .client .CoreV1 ().Services (service .GetNamespace ()).Delete (ctx , service .GetName (), v1.DeleteOptions {}); err != nil {
653
+ err := t .client .Do (ctx , func (ctx context.Context , client kubernetes.Interface ) error {
654
+ return client .CoreV1 ().Services (service .GetNamespace ()).Delete (ctx , service .GetName (), v1.DeleteOptions {})
655
+ })
656
+ if err != nil {
616
657
return err
617
658
}
618
659
// Wait for the service to disappear or for the context to expire.
619
660
for ctx .Err () == nil {
620
- if _ , err := t .client . CoreV1 (). Services ( service . GetNamespace ()). Get ( ctx , service . GetName (), v1. GetOptions {} ); err != nil {
661
+ if _ , err := t .GetService ( ctx , service ); err != nil {
621
662
return nil
622
663
}
623
664
select {
@@ -639,7 +680,7 @@ func (t *TestCluster) WaitForServiceReady(ctx context.Context, service *v13.Serv
639
680
case <- ctx .Done ():
640
681
return fmt .Errorf ("context expired waiting for service %q: %w (last: %v)" , service .GetName (), ctx .Err (), lastService )
641
682
case <- pollCh .C :
642
- s , err := t .client . CoreV1 (). Services ( service . GetNamespace ()). Get ( ctx , service . GetName (), v1. GetOptions {} )
683
+ s , err := t .GetService ( ctx , service )
643
684
if err != nil {
644
685
return fmt .Errorf ("cannot look up service %q: %w" , service .GetName (), err )
645
686
}
@@ -662,25 +703,40 @@ func (t *TestCluster) CreatePersistentVolume(ctx context.Context, volume *v13.Pe
662
703
if volume .GetObjectMeta ().GetNamespace () == "" {
663
704
volume .SetNamespace (NamespaceDefault )
664
705
}
665
- return t .client .CoreV1 ().PersistentVolumeClaims (volume .GetNamespace ()).Create (ctx , volume , v1.CreateOptions {})
706
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.PersistentVolumeClaim , error ) {
707
+ return client .CoreV1 ().PersistentVolumeClaims (volume .GetNamespace ()).Create (ctx , volume , v1.CreateOptions {})
708
+ })
666
709
}
667
710
668
711
// DeletePersistentVolume deletes a persistent volume.
669
712
func (t * TestCluster ) DeletePersistentVolume (ctx context.Context , volume * v13.PersistentVolumeClaim ) error {
670
- return t .client .CoreV1 ().PersistentVolumeClaims (volume .GetNamespace ()).Delete (ctx , volume .GetName (), v1.DeleteOptions {})
713
+ return t .client .Do (ctx , func (ctx context.Context , client kubernetes.Interface ) error {
714
+ return client .CoreV1 ().PersistentVolumeClaims (volume .GetNamespace ()).Delete (ctx , volume .GetName (), v1.DeleteOptions {})
715
+ })
671
716
}
672
717
673
718
// CreateDaemonset creates a daemonset with default options.
674
719
func (t * TestCluster ) CreateDaemonset (ctx context.Context , ds * appsv1.DaemonSet ) (* appsv1.DaemonSet , error ) {
675
720
if ds .GetObjectMeta ().GetNamespace () == "" {
676
721
ds .SetNamespace (NamespaceDefault )
677
722
}
678
- return t .client .AppsV1 ().DaemonSets (ds .GetNamespace ()).Create (ctx , ds , v1.CreateOptions {})
723
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* appsv1.DaemonSet , error ) {
724
+ return client .AppsV1 ().DaemonSets (ds .GetNamespace ()).Create (ctx , ds , v1.CreateOptions {})
725
+ })
726
+ }
727
+
728
+ // GetDaemonset gets a daemonset.
729
+ func (t * TestCluster ) GetDaemonset (ctx context.Context , ds * appsv1.DaemonSet ) (* appsv1.DaemonSet , error ) {
730
+ return request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* appsv1.DaemonSet , error ) {
731
+ return client .AppsV1 ().DaemonSets (ds .GetNamespace ()).Get (ctx , ds .GetName (), v1.GetOptions {})
732
+ })
679
733
}
680
734
681
735
// DeleteDaemonset deletes a daemonset from this cluster.
682
736
func (t * TestCluster ) DeleteDaemonset (ctx context.Context , ds * appsv1.DaemonSet ) error {
683
- return t .client .AppsV1 ().DaemonSets (ds .GetNamespace ()).Delete (ctx , ds .GetName (), v1.DeleteOptions {})
737
+ return t .client .Do (ctx , func (ctx context.Context , client kubernetes.Interface ) error {
738
+ return client .AppsV1 ().DaemonSets (ds .GetNamespace ()).Delete (ctx , ds .GetName (), v1.DeleteOptions {})
739
+ })
684
740
}
685
741
686
742
// GetPodsInDaemonSet returns the list of pods of the given DaemonSet.
@@ -689,7 +745,9 @@ func (t *TestCluster) GetPodsInDaemonSet(ctx context.Context, ds *appsv1.DaemonS
689
745
if appLabel , found := ds .Spec .Template .Labels [k8sApp ]; found {
690
746
listOptions .LabelSelector = fmt .Sprintf ("%s=%s" , k8sApp , appLabel )
691
747
}
692
- pods , err := t .client .CoreV1 ().Pods (ds .ObjectMeta .Namespace ).List (ctx , listOptions )
748
+ pods , err := request (ctx , t .client , func (ctx context.Context , client kubernetes.Interface ) (* v13.PodList , error ) {
749
+ return client .CoreV1 ().Pods (ds .ObjectMeta .Namespace ).List (ctx , listOptions )
750
+ })
693
751
if err != nil {
694
752
return nil , err
695
753
}
@@ -709,7 +767,7 @@ func (t *TestCluster) WaitForDaemonset(ctx context.Context, ds *appsv1.DaemonSet
709
767
defer pollCh .Stop ()
710
768
// Poll-based loop to wait for the DaemonSet to be ready.
711
769
for {
712
- d , err := t .client . AppsV1 (). DaemonSets ( ds . GetNamespace ()). Get ( ctx , ds . GetName (), v1. GetOptions {} )
770
+ d , err := t .GetDaemonset ( ctx , ds )
713
771
if err != nil {
714
772
return fmt .Errorf ("failed to get daemonset %q: %v" , ds .GetName (), err )
715
773
}
@@ -778,7 +836,7 @@ func (t *TestCluster) StreamDaemonSetLogs(ctx context.Context, ds *appsv1.Daemon
778
836
if _ , seen := nodesSeen [pod .Spec .NodeName ]; seen {
779
837
continue // Node already seen.
780
838
}
781
- logReader , err := t .client . CoreV1 (). Pods ( pod . GetNamespace ()). GetLogs ( pod . GetName () , & opts ). Stream ( ctx )
839
+ logReader , err := t .GetLogReader ( ctx , & pod , opts )
782
840
if err != nil {
783
841
// This can happen if the container hasn't run yet, for example
784
842
// because other init containers that run earlier are still executing.
@@ -813,7 +871,7 @@ Outer:
813
871
}
814
872
break Outer
815
873
case <- timeTicker .C :
816
- d , err := t .client . AppsV1 (). DaemonSets ( ds . GetNamespace ()). Get ( ctx , ds . GetName (), v1. GetOptions {} )
874
+ d , err := t .GetDaemonset ( ctx , ds )
817
875
if err != nil {
818
876
loopError = fmt .Errorf ("failed to get DaemonSet: %v" , err )
819
877
break Outer
0 commit comments