Prepare for clockwork v0.5.0 (#51292)

* Introduce the clocki package * Do root module fixes so e/ can be built * Use clocki in lib/backend tests * Genericfy test.Constructor and related methods
gravitational · Jan 23, 2025 · 1f7b09f · 1f7b09f
1 parent e04b9fb
commit 1f7b09f
Show file tree

Hide file tree

Showing 20 changed files with 145 additions and 49 deletions.
diff --git a/lib/auth/helpers_mfa.go b/lib/auth/helpers_mfa.go
@@ -33,6 +33,7 @@ import (
 	"github.com/gravitational/teleport/api/types"
 	"github.com/gravitational/teleport/lib/auth/mocku2f"
 	wantypes "github.com/gravitational/teleport/lib/auth/webauthntypes"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // TestDevice is a test MFA device.
@@ -190,9 +191,8 @@ func (d *TestDevice) solveAuthnTOTP(c *proto.MFAAuthenticateChallenge) (*proto.M
 	if d.clock == nil {
 		return nil, trace.BadParameter("clock not set")
 	}
-	if c, ok := d.clock.(clockwork.FakeClock); ok {
-		c.Advance(30 * time.Second)
-	}
+	clocki.Advance(d.clock, 30*time.Second)
+
 	code, err := totp.GenerateCode(d.TOTPSecret, d.clock.Now())
 	if err != nil {
 		return nil, trace.Wrap(err)
@@ -244,9 +244,7 @@ func (d *TestDevice) solveRegisterTOTP(c *proto.MFARegisterChallenge) (*proto.MF
 	if d.clock == nil {
 		return nil, trace.BadParameter("clock not set")
 	}
-	if c, ok := d.clock.(clockwork.FakeClock); ok {
-		c.Advance(30 * time.Second)
-	}
+	clocki.Advance(d.clock, 30*time.Second)
 
 	if c.GetTOTP().Algorithm != otp.AlgorithmSHA1.String() {
 		return nil, trace.BadParameter("unexpected TOTP challenge algorithm: %s", c.GetTOTP().Algorithm)

diff --git a/lib/backend/dynamo/atomicwrite_test.go b/lib/backend/dynamo/atomicwrite_test.go
@@ -28,9 +28,10 @@ import (
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
-func newAtomicWriteBackend(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteBackend(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 	dynamoCfg := map[string]interface{}{
 		"table_name":         dynamoDBTestTable(),
 		"poll_stream_period": 300 * time.Millisecond,

diff --git a/lib/backend/dynamo/dynamodbbk_test.go b/lib/backend/dynamo/dynamodbbk_test.go
@@ -43,6 +43,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 func TestMain(m *testing.M) {
@@ -74,7 +75,7 @@ func TestDynamoDB(t *testing.T) {
 		"poll_stream_period": 300 * time.Millisecond,
 	}
 
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		testCfg, err := test.ApplyOptions(options)
 		if err != nil {
 			return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/etcdbk/atomicwrite_test.go b/lib/backend/etcdbk/atomicwrite_test.go
@@ -23,15 +23,15 @@ import (
 	"testing"
 
 	"github.com/gravitational/trace"
-	"github.com/jonboulle/clockwork"
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // newAtomicWriteTestBackend builds a backend suitable for the atomic write test suite. Once all backends implement AtomicWrite,
 // it will be integrated into the main backend interface and we can get rid of this separate helper.
-func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 	opts, err := test.ApplyOptions(options)
 	if err != nil {
 		return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/etcdbk/etcd_test.go b/lib/backend/etcdbk/etcd_test.go
@@ -28,13 +28,13 @@ import (
 	"time"
 
 	"github.com/gravitational/trace"
-	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 const (
@@ -65,7 +65,7 @@ func TestEtcd(t *testing.T) {
 		t.Skip("This test requires etcd, run `make run-etcd` and set TELEPORT_ETCD_TEST=yes in your environment")
 	}
 
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		opts, err := test.ApplyOptions(options)
 		if err != nil {
 			return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/firestore/atomicwrite_test.go b/lib/backend/firestore/atomicwrite_test.go
@@ -27,9 +27,10 @@ import (
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
-func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 	cfg := firestoreParams()
 
 	testCfg, err := test.ApplyOptions(options)

diff --git a/lib/backend/firestore/firestorebk_test.go b/lib/backend/firestore/firestorebk_test.go
@@ -53,6 +53,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 func TestMain(m *testing.M) {
@@ -126,7 +127,7 @@ func TestFirestoreDB(t *testing.T) {
 	ensureTestsEnabled(t)
 	ensureEmulatorRunning(t, cfg)
 
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		testCfg, err := test.ApplyOptions(options)
 		if err != nil {
 			return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/lite/atomicwrite_test.go b/lib/backend/lite/atomicwrite_test.go
@@ -28,12 +28,13 @@ import (
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // newAtomicWriteTestBackendBuilder builds a backend suitable for the atomic write test suite. Once all backends implement AtomicWrite,
 // it will be integrated into the main backend interface and we can get rid of this separate helper.
-func newAtomicWriteTestBackendBuilder(t *testing.T) func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
-	return func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteTestBackendBuilder(t *testing.T) test.Constructor[clocki.FakeClock] {
+	return func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		clock := clockwork.NewFakeClock()
 
 		cfg, err := test.ApplyOptions(options)

diff --git a/lib/backend/lite/lite_test.go b/lib/backend/lite/lite_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 func TestMain(m *testing.M) {
@@ -39,7 +40,7 @@ func TestMain(m *testing.M) {
 }
 
 func TestLite(t *testing.T) {
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		clock := clockwork.NewFakeClock()
 
 		cfg, err := test.ApplyOptions(options)

diff --git a/lib/backend/memory/atomicwrite_test.go b/lib/backend/memory/atomicwrite_test.go
@@ -26,11 +26,12 @@ import (
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // newAtomicWriteTestBackend builds a backend suitable for the atomic write test suite. Once all backends implement AtomicWrite,
 // it will be integrated into the main backend interface and we can get rid of this separate helper.
-func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 	cfg, err := test.ApplyOptions(options)
 
 	if err != nil {

diff --git a/lib/backend/memory/memory_test.go b/lib/backend/memory/memory_test.go
@@ -32,6 +32,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 func TestMain(m *testing.M) {
@@ -40,7 +41,7 @@ func TestMain(m *testing.M) {
 }
 
 func TestMemory(t *testing.T) {
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		cfg, err := test.ApplyOptions(options)
 
 		if err != nil {

diff --git a/lib/backend/pgbk/atomicwrite_test.go b/lib/backend/pgbk/atomicwrite_test.go
@@ -29,6 +29,7 @@ import (
 
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // Testing requires a local psql backend to be set up, and for params to be passed via env. Ex:
@@ -37,7 +38,7 @@ import (
 
 // newAtomicWriteTestBackend builds a backend suitable for the atomic write test suite. Once all backends implement AtomicWrite,
 // it will be integrated into the main backend interface and we can get rid of this separate helper.
-func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func newAtomicWriteTestBackend(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 	testCfg, err := test.ApplyOptions(options)
 	if err != nil {
 		return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/pgbk/pgbk_test.go b/lib/backend/pgbk/pgbk_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend"
 	"github.com/gravitational/teleport/lib/backend/test"
 	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 func TestMain(m *testing.M) {
@@ -47,7 +48,7 @@ func TestPostgresBackend(t *testing.T) {
 		t.Skip("Postgres backend tests are disabled. Enable them by setting the TELEPORT_PGBK_TEST_PARAMS_JSON variable.")
 	}
 
-	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		testCfg, err := test.ApplyOptions(options)
 		if err != nil {
 			return nil, nil, trace.Wrap(err)

diff --git a/lib/backend/test/atomicwrite.go b/lib/backend/test/atomicwrite.go
@@ -30,9 +30,10 @@ import (
 	"golang.org/x/sync/errgroup"
 
 	"github.com/gravitational/teleport/lib/backend"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
-func RunAtomicWriteComplianceSuite(t *testing.T, newBackend Constructor) {
+func RunAtomicWriteComplianceSuite[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	t.Run("Move", func(t *testing.T) {
 		testAtomicWriteMove(t, newBackend)
 	})
@@ -59,7 +60,7 @@ func RunAtomicWriteComplianceSuite(t *testing.T, newBackend Constructor) {
 }
 
 // testAtomicWriteMove verifies the correct behavior of "move" operations.
-func testAtomicWriteMove(t *testing.T, newBackend Constructor) {
+func testAtomicWriteMove[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	bk, _, err := newBackend()
 	require.NoError(t, err)
 
@@ -121,7 +122,7 @@ func testAtomicWriteMove(t *testing.T, newBackend Constructor) {
 
 // testAtomicWriteLock verifies correct behavior of various "lock" patterns (i.e. where some update on key X is conditional on
 // the state of key Y).
-func testAtomicWriteLock(t *testing.T, newBackend Constructor) {
+func testAtomicWriteLock[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	bk, _, err := newBackend()
 	require.NoError(t, err)
 
@@ -275,7 +276,7 @@ func testAtomicWriteLock(t *testing.T, newBackend Constructor) {
 }
 
 // testAtomicWriteMax verifies correct behavior of very large atomic writes.
-func testAtomicWriteMax(t *testing.T, newBackend Constructor) {
+func testAtomicWriteMax[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	bk, _, err := newBackend()
 	require.NoError(t, err)
 
@@ -361,7 +362,7 @@ func testAtomicWriteMax(t *testing.T, newBackend Constructor) {
 }
 
 // testAtomicWriteConcurrent is a sanity-check intended to verify the correctness of AtomicWrite under high concurrency.
-func testAtomicWriteConcurrent(t *testing.T, newBackend Constructor) {
+func testAtomicWriteConcurrent[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	const (
 		increments = 200
 		workers    = 20
@@ -454,7 +455,7 @@ func testAtomicWriteConcurrent(t *testing.T, newBackend Constructor) {
 // testAtomicWriteNonConflicting verifies that non-conflicting but overlapping transactions all succeed
 // on the first attempt when running concurrently, meaning that backends that treat overlap as conflict (e.g. dynamo)
 // handle such conflicts internally.
-func testAtomicWriteNonConflicting(t *testing.T, newBackend Constructor) {
+func testAtomicWriteNonConflicting[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	const workers = 60
 
 	bk, _, err := newBackend()
@@ -521,7 +522,7 @@ func testAtomicWriteNonConflicting(t *testing.T, newBackend Constructor) {
 // testAtomicWriteOther verifies some minor edge-cases that may not be covered by other tests. Specifically,
 // it verifies that Item.Key has no effect on writes or subsequent reads, and that ineffectual writes still
 // update the value of revision.
-func testAtomicWriteOther(t *testing.T, newBackend Constructor) {
+func testAtomicWriteOther[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
 	bk, _, err := newBackend()
 	require.NoError(t, err)
 

diff --git a/lib/backend/test/atomicwrite_shim.go b/lib/backend/test/atomicwrite_shim.go
@@ -26,17 +26,17 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/gravitational/trace"
-	"github.com/jonboulle/clockwork"
 
 	"github.com/gravitational/teleport/lib/backend"
+	"github.com/gravitational/teleport/lib/utils/clocki"
 )
 
 // RunBackendComplianceSuiteWithAtomicWriteShim runs the old backend compliance suite against the provided backend
 // with a shim that converts all calls to single-write methods (all write methods but DeleteRange) into calls to
 // AtomicWrite. This is done to ensure that the relationship between the conditional actions of AtomicWrite and the
 // single-write methods is well defined, and to improve overall coverage of AtomicWrite implementations via reuse.
-func RunBackendComplianceSuiteWithAtomicWriteShim(t *testing.T, newBackend Constructor) {
-	RunBackendComplianceSuite(t, func(options ...ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+func RunBackendComplianceSuiteWithAtomicWriteShim[FakeClock clocki.FakeClock](t *testing.T, newBackend Constructor[FakeClock]) {
+	RunBackendComplianceSuite(t, func(options ...ConstructionOption) (backend.Backend, clocki.FakeClock, error) {
 		bk, clock, err := newBackend(options...)
 		if err != nil {
 			return nil, nil, trace.Wrap(err)