mockClient.go

Documentation: edge-infra.dev/pkg/edge/device-registrar/utils/test

     1  package test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
     8  	"k8s.io/apimachinery/pkg/runtime"
     9  	"k8s.io/apimachinery/pkg/runtime/schema"
    10  	"k8s.io/apimachinery/pkg/util/sets"
    11  	"sigs.k8s.io/controller-runtime/pkg/client"
    12  	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
    13  	"sigs.k8s.io/controller-runtime/pkg/client/fake"
    14  )
    15  
    16  // ErrorInjectingFakeClient is a fake client that can be used to inject errors into the K8S client
    17  type ErrorInjectingFakeClient struct {
    18  	client.WithWatch
    19  	withSubResourceSimulation bool
    20  	CustomSubResourceWriter   *CustomSubResourceWriter
    21  	withObjects               []client.Object
    22  	withStatusSubresource     []client.Object
    23  	failOnCreate              bool
    24  	failOnUpdate              bool
    25  	failOnList                bool
    26  	failOnDelete              bool
    27  	failOnGet                 bool
    28  	failOnStatus              bool
    29  }
    30  
    31  func (c *ErrorInjectingFakeClient) Status() client.SubResourceWriter {
    32  	if c.failOnStatus {
    33  		// use the custom subresource writer to inject errors
    34  		if c.CustomSubResourceWriter == nil {
    35  			return c.NewCustomSubResourceWriter()
    36  		}
    37  		return c.CustomSubResourceWriter
    38  	}
    39  	return c.WithWatch.Status()
    40  }
    41  
    42  func NewErrorInjectingFakeClient(scheme *runtime.Scheme, withSubResourceSimulation bool, initObjects ...client.Object) *ErrorInjectingFakeClient {
    43  	builder := fake.NewClientBuilder().
    44  		WithScheme(scheme).
    45  		WithObjects(initObjects...).
    46  		WithStatusSubresource(initObjects...)
    47  	return &ErrorInjectingFakeClient{
    48  		withSubResourceSimulation: withSubResourceSimulation,
    49  		WithWatch:                 builder.Build(),
    50  		withObjects:               initObjects,
    51  		withStatusSubresource:     initObjects,
    52  	}
    53  }
    54  
    55  // Copied from sigs.k8s.io/controller-runtime@v0.16.3/pkg/client/fake/client.go#1217
    56  var inTreeResourcesWithStatus = []schema.GroupVersionKind{
    57  	{Version: "v1", Kind: "Namespace"},
    58  	{Version: "v1", Kind: "Node"},
    59  	{Version: "v1", Kind: "PersistentVolumeClaim"},
    60  	{Version: "v1", Kind: "PersistentVolume"},
    61  	{Version: "v1", Kind: "Pod"},
    62  	{Version: "v1", Kind: "ReplicationController"},
    63  	{Version: "v1", Kind: "Service"},
    64  
    65  	{Group: "apps", Version: "v1", Kind: "Deployment"},
    66  	{Group: "apps", Version: "v1", Kind: "DaemonSet"},
    67  	{Group: "apps", Version: "v1", Kind: "ReplicaSet"},
    68  	{Group: "apps", Version: "v1", Kind: "StatefulSet"},
    69  
    70  	{Group: "autoscaling", Version: "v1", Kind: "HorizontalPodAutoscaler"},
    71  
    72  	{Group: "batch", Version: "v1", Kind: "CronJob"},
    73  	{Group: "batch", Version: "v1", Kind: "Job"},
    74  
    75  	{Group: "certificates.k8s.io", Version: "v1", Kind: "CertificateSigningRequest"},
    76  
    77  	{Group: "networking.k8s.io", Version: "v1", Kind: "Ingress"},
    78  	{Group: "networking.k8s.io", Version: "v1", Kind: "NetworkPolicy"},
    79  
    80  	{Group: "policy", Version: "v1", Kind: "PodDisruptionBudget"},
    81  
    82  	{Group: "storage.k8s.io", Version: "v1", Kind: "VolumeAttachment"},
    83  
    84  	{Group: "apiextensions.k8s.io", Version: "v1", Kind: "CustomResourceDefinition"},
    85  
    86  	{Group: "flowcontrol.apiserver.k8s.io", Version: "v1beta2", Kind: "FlowSchema"},
    87  	{Group: "flowcontrol.apiserver.k8s.io", Version: "v1beta2", Kind: "PriorityLevelConfiguration"},
    88  }
    89  
    90  func (c *ErrorInjectingFakeClient) Create(ctx context.Context, obj client.Object, opts ...client.CreateOption) error {
    91  	if c.failOnCreate {
    92  		return errors.New("create error")
    93  	}
    94  
    95  	if err := c.WithWatch.Create(ctx, obj, opts...); err != nil {
    96  		return err
    97  	}
    98  	if !c.withSubResourceSimulation {
    99  		return nil
   100  	}
   101  
   102  	c.withObjects = append(c.withObjects, obj)
   103  	err := c.recreateFake(ctx)
   104  	if err != nil {
   105  		return err
   106  	}
   107  
   108  	return nil
   109  }
   110  
   111  func (c *ErrorInjectingFakeClient) Delete(ctx context.Context, obj client.Object, opts ...client.DeleteOption) error {
   112  	if c.failOnDelete {
   113  		return errors.New("delete error")
   114  	}
   115  
   116  	if err := c.WithWatch.Delete(ctx, obj, opts...); err != nil {
   117  		return err
   118  	}
   119  
   120  	return nil
   121  }
   122  
   123  func (c *ErrorInjectingFakeClient) Update(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
   124  	if c.failOnUpdate {
   125  		return errors.New("update error")
   126  	}
   127  	return c.WithWatch.Update(ctx, obj, opts...)
   128  }
   129  
   130  // create func for List method
   131  func (c *ErrorInjectingFakeClient) List(ctx context.Context, list client.ObjectList, opts ...client.ListOption) error {
   132  	if c.failOnList {
   133  		return errors.New("unsupported list type")
   134  	}
   135  	return c.WithWatch.List(ctx, list, opts...)
   136  }
   137  
   138  // create func for Get method
   139  func (c *ErrorInjectingFakeClient) Get(ctx context.Context, key client.ObjectKey, obj client.Object, opts ...client.GetOption) error {
   140  	if c.failOnGet {
   141  		return errors.New("unsupported get type")
   142  	}
   143  	return c.WithWatch.Get(ctx, key, obj, opts...)
   144  }
   145  
   146  func (c *ErrorInjectingFakeClient) recreateFake(ctx context.Context) error {
   147  	// This is a pretty disgusting hack to get around the fact that the fake client doesn't support creating status subresources
   148  	// Discussed here https://github.com/kubernetes-sigs/controller-runtime/issues/2386#issuecomment-1607768830
   149  	// and here https://github.com/kubernetes-sigs/controller-runtime/issues/2362#issuecomment-1699415588
   150  	// and here https://stackoverflow.com/questions/77489441/go-k8s-controller-runtime-upgrade-fake-client-lacks-functionality
   151  
   152  	// We need to collect all the objects from the fake client, and then re-create the fake client with the status subresource for the created object
   153  
   154  	gvks := sets.New(inTreeResourcesWithStatus...)
   155  	for _, o := range c.withObjects {
   156  		gvk, err := apiutil.GVKForObject(o, c.Scheme())
   157  		if err != nil {
   158  			return err
   159  		}
   160  		gvks.Insert(gvk)
   161  	}
   162  
   163  	var objs []client.Object
   164  	for _, gvk := range gvks.UnsortedList() {
   165  		objList := &unstructured.UnstructuredList{}
   166  		objList.SetGroupVersionKind(schema.GroupVersionKind{Group: gvk.Group, Version: gvk.Version, Kind: gvk.Kind})
   167  
   168  		err := c.List(ctx, objList)
   169  		if err != nil {
   170  			return err
   171  		}
   172  
   173  		for _, o := range objList.Items {
   174  			objs = append(objs, o.DeepCopy())
   175  		}
   176  	}
   177  
   178  	initObjs := sets.New(objs...).UnsortedList()
   179  	c.withObjects = initObjs
   180  	c.withStatusSubresource = initObjs
   181  
   182  	c.WithWatch = fake.NewClientBuilder().
   183  		WithScheme(c.Scheme()).
   184  		WithObjects(c.withObjects...).
   185  		WithStatusSubresource(c.withStatusSubresource...).
   186  		Build()
   187  	return nil
   188  }
   189  
   190  // SetFailOnCreate sets the flag to fail on create
   191  func (c *ErrorInjectingFakeClient) SetFailOnCreate() {
   192  	c.failOnCreate = true
   193  }
   194  
   195  // SetFailOnUpdate sets the flag to fail on update
   196  func (c *ErrorInjectingFakeClient) SetFailOnUpdate() {
   197  	c.failOnUpdate = true
   198  }
   199  
   200  // SetFailOnList sets the flag to fail on list
   201  func (c *ErrorInjectingFakeClient) SetFailOnList() {
   202  	c.failOnList = true
   203  }
   204  
   205  // SetFailOnDelete sets the flag to fail on delete
   206  func (c *ErrorInjectingFakeClient) SetFailOnDelete() {
   207  	c.failOnDelete = true
   208  }
   209  
   210  // SetFailOnGet sets the flag to fail on get
   211  func (c *ErrorInjectingFakeClient) SetFailOnGet() {
   212  	c.failOnGet = true
   213  }
   214  
   215  // SetFailOnStatus sets the flag to fail on status
   216  func (c *ErrorInjectingFakeClient) SetFailOnStatus() {
   217  	c.failOnStatus = true
   218  }
   219  
   220  // CustomSubResourceWriter is a custom implementation of the client.SubResourceWriter interface
   221  type CustomSubResourceWriter struct {
   222  	client       client.Client
   223  	failOnUpdate bool
   224  	failOnCreate bool
   225  	failOnPatch  bool
   226  }
   227  
   228  // setFailOnUpdate sets the flag to fail on update
   229  func (w *CustomSubResourceWriter) SetFailOnUpdate() {
   230  	w.failOnUpdate = true
   231  }
   232  
   233  // setFailOnCreate sets the flag to fail on create
   234  func (w *CustomSubResourceWriter) SetFailOnCreate() {
   235  	w.failOnCreate = true
   236  }
   237  
   238  // setFailOnPatch sets the flag to fail on patch
   239  func (w *CustomSubResourceWriter) SetFailOnPatch() {
   240  	w.failOnPatch = true
   241  }
   242  
   243  func (w *CustomSubResourceWriter) Create(ctx context.Context, obj client.Object, subResource client.Object, opts ...client.SubResourceCreateOption) error {
   244  	if w.failOnCreate {
   245  		return errors.New("create error")
   246  	}
   247  	return w.client.SubResource(obj.GetResourceVersion()).Create(ctx, obj, subResource, opts...)
   248  }
   249  
   250  func (w *CustomSubResourceWriter) Update(ctx context.Context, obj client.Object, opts ...client.SubResourceUpdateOption) error {
   251  	if w.failOnUpdate {
   252  		return errors.New("update error")
   253  	}
   254  	return w.client.SubResource(obj.GetResourceVersion()).Update(ctx, obj, opts...)
   255  }
   256  
   257  func (w *CustomSubResourceWriter) Patch(ctx context.Context, obj client.Object, patch client.Patch, opts ...client.SubResourcePatchOption) error {
   258  	if w.failOnPatch {
   259  		return errors.New("patch error")
   260  	}
   261  	return w.client.SubResource(obj.GetResourceVersion()).Patch(ctx, obj, patch, opts...)
   262  }
   263  
   264  func (c *ErrorInjectingFakeClient) NewCustomSubResourceWriter() *CustomSubResourceWriter {
   265  	c.CustomSubResourceWriter = &CustomSubResourceWriter{
   266  		client: c.WithWatch,
   267  	}
   268  	return c.CustomSubResourceWriter
   269  }
   270
View as plain text