kpoll.go

Documentation: edge-infra.dev/test/f2/x/ktest/kpoll

     1  package kpoll
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"gotest.tools/v3/assert/cmp"
     9  	"gotest.tools/v3/poll"
    10  	"k8s.io/apimachinery/pkg/api/errors"
    11  	"sigs.k8s.io/controller-runtime/pkg/client"
    12  
    13  	"edge-infra.dev/pkg/k8s/runtime/inventory"
    14  	"edge-infra.dev/pkg/k8s/testing/kmp"
    15  )
    16  
    17  // KPoll allows asserting K8s Objects against a live cluster's state with
    18  // default polling options and provides access to common assertions for K8s
    19  // objects.
    20  //
    21  // KPoll should be treated like other assertion helpers: instantiated for each
    22  // individual test.
    23  //
    24  // When instantiated as part of a test hook, this ensures that all assertions
    25  // use a consistent and configurable timeout by default.
    26  type KPoll struct {
    27  	c       client.Client
    28  	ctx     context.Context
    29  	timeout time.Duration
    30  	tick    time.Duration
    31  }
    32  
    33  func New(ctx context.Context, c client.Client, timeout, tick time.Duration) *KPoll {
    34  	return &KPoll{c, ctx, timeout, tick}
    35  }
    36  
    37  // WaitOn wraps [gotest.tools/v3/poll.WaitOn] with default options based on the
    38  // KPoll configuration.
    39  func (k *KPoll) WaitOn(t *testing.T, c poll.Check, opts ...poll.SettingOp) {
    40  	t.Helper()
    41  	poll.WaitOn(t, c, append(k.pollSettings(), opts...)...)
    42  }
    43  
    44  // Compare returns a [gotest.tools/v3/poll.Result] by evaluating Object o against
    45  // Komparison c. The Object is re-fetched on each polling loop, allowing a
    46  // standard [edge-infra.dev/pkg/k8s/testing/kmp.Komparison] to be continually
    47  // evaluated against a cluster's state.
    48  //
    49  // The poll.Result return can be used to wrap a Komparison with additional logic
    50  // in an inline poll.Check:
    51  //
    52  //	kwait.On(t, func(t poll.LogT) poll.Result {
    53  //		// Evaluate check precondition
    54  //		// Perform comparison as polling step
    55  //		return kwait.Compare(ctx, obj, myComparison)
    56  //	})
    57  func (k *KPoll) Compare(o client.Object, c kmp.Komparison, options ...CheckOption) poll.Result {
    58  	opts := k.makeOpts(options...)
    59  	return poll.Compare(func() cmp.Result {
    60  		if opts.fetch {
    61  			err := opts.client.Get(opts.ctx, client.ObjectKeyFromObject(o), o)
    62  			if err != nil && !(errors.IsGone(err) || errors.IsNotFound(err)) {
    63  				return cmp.ResultFromError(err)
    64  			}
    65  		}
    66  		return c(o)
    67  	})
    68  }
    69  
    70  // CompareAll is [KPoll.Compare] for sets of objects. The comparison results are
    71  // aggregated to produce the overall result.
    72  func (k *KPoll) CompareAll(objs []client.Object, c kmp.Komparison, options ...CheckOption) poll.Result {
    73  	opts := k.makeOpts(options...)
    74  	results := make([]poll.Result, len(objs), 0)
    75  	for i, o := range objs {
    76  		if opts.fetch {
    77  			err := opts.client.Get(opts.ctx, client.ObjectKeyFromObject(o), o)
    78  			if err != nil && !(errors.IsGone(err) || errors.IsNotFound(err)) {
    79  				return poll.Error(err)
    80  			}
    81  		}
    82  		results[i] = poll.Compare(func() cmp.Result { return c(o) })
    83  	}
    84  
    85  	return JoinResults(results...)
    86  }
    87  
    88  // Check creates a [gotest.tools/v3/poll.Check] which evaluates Object o using
    89  // Komparison c on each interval of a [KPoll.WaitOn] call. It uses [KPoll.Compare]
    90  // to wrap the Komparison into a function that produces a polling result so that
    91  // the check is always executed against the cluster's live state.
    92  func (k *KPoll) Check(o client.Object, c kmp.Komparison, options ...CheckOption) poll.Check {
    93  	return func(_ poll.LogT) poll.Result {
    94  		return k.Compare(o, c, options...)
    95  	}
    96  }
    97  
    98  // CheckAll is [KPoll.Check] for sets of objects.
    99  func (k *KPoll) CheckAll(objs []client.Object, c kmp.Komparison, options ...CheckOption) poll.Check {
   100  	return func(_ poll.LogT) poll.Result {
   101  		return k.CompareAll(objs, c, options...)
   102  	}
   103  }
   104  
   105  // ObjExists checks that a single object exists on the cluster.
   106  //
   107  // [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
   108  // existence.
   109  func (k *KPoll) ObjExists(o client.Object, options ...CheckOption) poll.Check {
   110  	opts := k.makeOpts(options...)
   111  	return Exists(opts.ctx, opts.client, o)
   112  }
   113  
   114  // ObjsExist checks that all objects exist on the cluster.
   115  //
   116  // [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
   117  // existence.
   118  func (k *KPoll) ObjsExist(objs []client.Object, options ...CheckOption) poll.Check {
   119  	opts := k.makeOpts(options...)
   120  	return Exists(opts.ctx, opts.client, objs...)
   121  }
   122  
   123  // ObjDeleted checks that a single object doesn't exist on the cluster.
   124  //
   125  // [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
   126  // existence.
   127  func (k *KPoll) ObjDeleted(o client.Object, options ...CheckOption) poll.Check {
   128  	opts := k.makeOpts(options...)
   129  	return Deleted(opts.ctx, opts.client, o)
   130  }
   131  
   132  // ObjsDeleted checks that none of the provided objects exist on the cluster.
   133  //
   134  // [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
   135  // existence.
   136  func (k *KPoll) ObjsDeleted(objs []client.Object, options ...CheckOption) poll.Check {
   137  	opts := k.makeOpts(options...)
   138  	return Deleted(opts.ctx, opts.client, objs...)
   139  }
   140  
   141  // InventoryExists checks that all member objects of the expected inventory
   142  // exist on the cluster.
   143  //
   144  // [edge-infra.dev/pkg/k8s/object.IsDeleted] is used to determine object
   145  // existence.
   146  func (k *KPoll) InventoryExists(exp *inventory.ResourceInventory, options ...CheckOption) poll.Check {
   147  	opts := k.makeOpts(options...)
   148  	return InventoryExists(opts.ctx, opts.client, exp)
   149  }
   150  
   151  func (k *KPoll) InventoryPruned(curr, old *inventory.ResourceInventory, options ...CheckOption) poll.Check {
   152  	opts := k.makeOpts(options...)
   153  	return InventoryPruned(opts.ctx, opts.client, curr, old)
   154  }
   155  
   156  func (k *KPoll) makeOpts(options ...CheckOption) *checkOpts {
   157  	opts := &checkOpts{
   158  		ctx:    k.ctx,
   159  		client: k.c,
   160  		fetch:  true,
   161  	}
   162  	for _, option := range options {
   163  		option(opts)
   164  	}
   165  	return opts
   166  }
   167  
   168  // pollSettings returns the default [gotest.tools/v3/poll.WaitOn] options based
   169  // on the object's state.
   170  func (k *KPoll) pollSettings() []poll.SettingOp {
   171  	return []poll.SettingOp{poll.WithDelay(k.tick), poll.WithTimeout(k.timeout)}
   172  }
   173
View as plain text