diff --git a/injectors.go b/injectors.go new file mode 100644 index 0000000..604005a --- /dev/null +++ b/injectors.go @@ -0,0 +1,21 @@ +package ntest + +import ( + "context" +) + +// This file contains example injectors that may be useful + +// Cancel is the injected type for the function type that will cancel +// a Context that has been augmented with AutoCancel. +type Cancel func() + +// AutoCancel adjusts context.Context so that it will be cancelled +// when the test finishes. It can be cancelled early by calling +// the returned Cancel function. +func AutoCancel(ctx context.Context, t T) (context.Context, Cancel) { + ctx, cancel := context.WithCancel(ctx) + onlyOnce := onceFunc(cancel) + t.Cleanup(onlyOnce) + return ctx, onlyOnce +} diff --git a/injectors_test.go b/injectors_test.go new file mode 100644 index 0000000..7fbf4c2 --- /dev/null +++ b/injectors_test.go @@ -0,0 +1,29 @@ +package ntest_test + +import ( + "context" + "testing" + + "github.com/memsql/ntest" +) + +func TestCancel(t *testing.T) { + t.Parallel() + ntest.RunTest(t, context.Background, ntest.AutoCancel, func( + ctx context.Context, + cancel ntest.Cancel, + ) { + select { + case <-ctx.Done(): + t.Fatal("context is cancelled before cancel()") + default: + } + cancel() + select { + case <-ctx.Done(): + t.Log("context is cancelled") + default: + t.Fatal("context is not cancelled after cancel()") + } + }) +} diff --git a/once.go b/once.go new file mode 100644 index 0000000..545d8bd --- /dev/null +++ b/once.go @@ -0,0 +1,7 @@ +//go:build go1.21 + +package ntest + +import "sync" + +var onceFunc = sync.OnceFunc diff --git a/once_old.go b/once_old.go new file mode 100644 index 0000000..eeba032 --- /dev/null +++ b/once_old.go @@ -0,0 +1,12 @@ +//go:build !go1.21 + +package ntest + +import "sync" + +func onceFunc(f func()) func() { + var once sync.Once + return func() { + once.Do(f) + } +}