Update existing rapid tests to new version

consensus/hotstuff/cruisectl/transition_time_test.go

@@ -80,9 +80,9 @@ func TestParseTransition_Invalid(t *testing.T) {
 
 // drawTransitionTime draws a random EpochTransitionTime.
 func drawTransitionTime(t *rapid.T) EpochTransitionTime {
-	day := time.Weekday(rapid.IntRange(0, 6).Draw(t, "wd").(int))
-	hour := rapid.Uint8Range(0, 23).Draw(t, "h").(uint8)
-	minute := rapid.Uint8Range(0, 59).Draw(t, "m").(uint8)
+	day := time.Weekday(rapid.IntRange(0, 6).Draw(t, "wd"))
+	hour := rapid.Uint8Range(0, 23).Draw(t, "h")
+	minute := rapid.Uint8Range(0, 59).Draw(t, "m")
 	return EpochTransitionTime{day, hour, minute}
 }
 
@@ -111,8 +111,8 @@ func TestInferTargetEndTime_Fixture(t *testing.T) {
 func TestInferTargetEndTime_Rapid(t *testing.T) {
 	rapid.Check(t, func(t *rapid.T) {
 		ett := drawTransitionTime(t)
-		curTime := time.Unix(rapid.Int64().Draw(t, "ref_unix").(int64), 0).UTC()
-		epochFractionComplete := rapid.Float64Range(0, 1).Draw(t, "pct_complete").(float64)
+		curTime := time.Unix(rapid.Int64().Draw(t, "ref_unix"), 0).UTC()
+		epochFractionComplete := rapid.Float64Range(0, 1).Draw(t, "pct_complete")
 		epochFractionRemaining := 1.0 - epochFractionComplete
 
 		target := ett.inferTargetEndTime(curTime, epochFractionComplete)
@@ -131,7 +131,7 @@ func TestInferTargetEndTime_Rapid(t *testing.T) {
 func TestFindNearestTargetTime(t *testing.T) {
 	rapid.Check(t, func(t *rapid.T) {
 		ett := drawTransitionTime(t)
-		ref := time.Unix(rapid.Int64().Draw(t, "ref_unix").(int64), 0).UTC()
+		ref := time.Unix(rapid.Int64().Draw(t, "ref_unix"), 0).UTC()
 
 		nearest := ett.findNearestTargetTime(ref)
 		distance := nearest.Sub(ref).Abs()

consensus/hotstuff/votecollector/combined_vote_processor_v2_test.go

@@ -118,7 +118,7 @@ func (s *CombinedVoteProcessorV2TestSuite) TestProcess_InvalidSignatureFormat()
 	rapid.Check(s.T(), func(t *rapid.T) {
 		// create a signature with invalid length
 		vote := unittest.VoteForBlockFixture(s.proposal.Block, func(vote *model.Vote) {
-			vote.SigData = unittest.RandomBytes(generator.Draw(t, "sig-size").(int))
+			vote.SigData = unittest.RandomBytes(generator.Draw(t, "sig-size"))
 		})
 		err := s.processor.Process(vote)
 		require.Error(s.T(), err)
@@ -434,8 +434,8 @@ func TestCombinedVoteProcessorV2_PropertyCreatingQCCorrectness(testifyT *testing
 
 	rapid.Check(testifyT, func(t *rapid.T) {
 		// draw participants in range 1 <= participants <= maxParticipants
-		participants := rapid.Uint64Range(1, maxParticipants).Draw(t, "participants").(uint64)
-		beaconSignersCount := rapid.Uint64Range(participants/2+1, participants).Draw(t, "beaconSigners").(uint64)
+		participants := rapid.Uint64Range(1, maxParticipants).Draw(t, "participants")
+		beaconSignersCount := rapid.Uint64Range(participants/2+1, participants).Draw(t, "beaconSigners")
 		stakingSignersCount := participants - beaconSignersCount
 		require.Equal(t, participants, stakingSignersCount+beaconSignersCount)
 
@@ -638,20 +638,20 @@ func TestCombinedVoteProcessorV2_PropertyCreatingQCCorrectness(testifyT *testing
 func TestCombinedVoteProcessorV2_PropertyCreatingQCLiveness(testifyT *testing.T) {
 	rapid.Check(testifyT, func(t *rapid.T) {
 		// draw beacon signers in range 1 <= beaconSignersCount <= 53
-		beaconSignersCount := rapid.Uint64Range(1, 53).Draw(t, "beaconSigners").(uint64)
+		beaconSignersCount := rapid.Uint64Range(1, 53).Draw(t, "beaconSigners")
 		// draw staking signers in range 0 <= stakingSignersCount <= 10
-		stakingSignersCount := rapid.Uint64Range(0, 10).Draw(t, "stakingSigners").(uint64)
+		stakingSignersCount := rapid.Uint64Range(0, 10).Draw(t, "stakingSigners")
 
 		stakingWeightRange, beaconWeightRange := rapid.Uint64Range(1, 10), rapid.Uint64Range(1, 10)
 
 		minRequiredWeight := uint64(0)
 		// draw weight for each signer randomly
 		stakingSigners := unittest.IdentityListFixture(int(stakingSignersCount), func(identity *flow.Identity) {
-			identity.InitialWeight = stakingWeightRange.Draw(t, identity.String()).(uint64)
+			identity.InitialWeight = stakingWeightRange.Draw(t, identity.String())
 			minRequiredWeight += identity.InitialWeight
 		})
 		beaconSigners := unittest.IdentityListFixture(int(beaconSignersCount), func(identity *flow.Identity) {
-			identity.InitialWeight = beaconWeightRange.Draw(t, identity.String()).(uint64)
+			identity.InitialWeight = beaconWeightRange.Draw(t, identity.String())
 			minRequiredWeight += identity.InitialWeight
 		})
 

consensus/hotstuff/votecollector/combined_vote_processor_v3_test.go

@@ -434,8 +434,8 @@ func TestCombinedVoteProcessorV3_PropertyCreatingQCCorrectness(testifyT *testing
 
 	rapid.Check(testifyT, func(t *rapid.T) {
 		// draw participants in range 1 <= participants <= maxParticipants
-		participants := rapid.Uint64Range(1, maxParticipants).Draw(t, "participants").(uint64)
-		beaconSignersCount := rapid.Uint64Range(participants/2+1, participants).Draw(t, "beaconSigners").(uint64)
+		participants := rapid.Uint64Range(1, maxParticipants).Draw(t, "participants")
+		beaconSignersCount := rapid.Uint64Range(participants/2+1, participants).Draw(t, "beaconSigners")
 		stakingSignersCount := participants - beaconSignersCount
 		require.Equal(t, participants, stakingSignersCount+beaconSignersCount)
 
@@ -749,20 +749,20 @@ func TestCombinedVoteProcessorV3_OnlyRandomBeaconSigners(testifyT *testing.T) {
 func TestCombinedVoteProcessorV3_PropertyCreatingQCLiveness(testifyT *testing.T) {
 	rapid.Check(testifyT, func(t *rapid.T) {
 		// draw beacon signers in range 1 <= beaconSignersCount <= 53
-		beaconSignersCount := rapid.Uint64Range(1, 53).Draw(t, "beaconSigners").(uint64)
+		beaconSignersCount := rapid.Uint64Range(1, 53).Draw(t, "beaconSigners")
 		// draw staking signers in range 0 <= stakingSignersCount <= 10
-		stakingSignersCount := rapid.Uint64Range(0, 10).Draw(t, "stakingSigners").(uint64)
+		stakingSignersCount := rapid.Uint64Range(0, 10).Draw(t, "stakingSigners")
 
 		stakingWeightRange, beaconWeightRange := rapid.Uint64Range(1, 10), rapid.Uint64Range(1, 10)
 
 		minRequiredWeight := uint64(0)
 		// draw weight for each signer randomly
 		stakingSigners := unittest.IdentityListFixture(int(stakingSignersCount), func(identity *flow.Identity) {
-			identity.InitialWeight = stakingWeightRange.Draw(t, identity.String()).(uint64)
+			identity.InitialWeight = stakingWeightRange.Draw(t, identity.String())
 			minRequiredWeight += identity.InitialWeight
 		})
 		beaconSigners := unittest.IdentityListFixture(int(beaconSignersCount), func(identity *flow.Identity) {
-			identity.InitialWeight = beaconWeightRange.Draw(t, identity.String()).(uint64)
+			identity.InitialWeight = beaconWeightRange.Draw(t, identity.String())
 			minRequiredWeight += identity.InitialWeight
 		})
 

go.sum

@@ -263,7 +263,6 @@ github.com/cockroachdb/pebble v0.0.0-20230928194634-aa077af62593 h1:aPEJyR4rPBvD
 github.com/cockroachdb/pebble v0.0.0-20230928194634-aa077af62593/go.mod h1:6hk1eMY/u5t+Cf18q5lFMUA1Rc+Sm5I6Ra1QuPyxXCo=
 github.com/cockroachdb/redact v1.1.3 h1:AKZds10rFSIj7qADf0g46UixK8NNLwWTNdCIGS5wfSQ=
 github.com/cockroachdb/redact v1.1.3/go.mod h1:BVNblN9mBWFyMyqK1k3AAiSxhvhfK2oOZZ2lK+dpvRg=
-github.com/cockroachdb/sentry-go v0.6.1-cockroachdb.2/go.mod h1:8BT+cPK6xvFOcRlk0R8eg+OTkcqI6baNH4xAkpiYVvQ=
 github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 h1:zuQyyAKVxetITBuuhv3BI9cMrmStnpT18zmgmTxunpo=
 github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06/go.mod h1:7nc4anLGjupUW/PeY5qiNYsdNXj7zopG+eqsS7To5IQ=
 github.com/codahale/hdrhistogram v0.0.0-20161010025455-3a0bb77429bd/go.mod h1:sE/e/2PUdi/liOCUjSTXgM1o87ZssimdTWN964YiIeI=
@@ -858,7 +857,6 @@ github.com/json-iterator/go v1.1.8/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/u
 github.com/json-iterator/go v1.1.9/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
-github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU=
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
 github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
@@ -1244,7 +1242,6 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJ
 github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
-github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/montanaflynn/stats v0.6.6 h1:Duep6KMIDpY4Yo11iFsvyqJDyfzLF9+sndUKT+v64GQ=
 github.com/montanaflynn/stats v0.6.6/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
 github.com/moul/http2curl v1.0.0/go.mod h1:8UbvGypXm98wA/IqH45anm5Y2Z6ep6O31QGOAZ3H0fQ=
@@ -2478,7 +2475,6 @@ lukechampine.com/blake3 v1.2.1/go.mod h1:0OFRp7fBtAylGVCO40o87sbupkyIGgbpv1+M1k1
 nhooyr.io/websocket v1.8.6/go.mod h1:B70DZP8IakI65RVQ51MsWP/8jndNma26DVA/nFSCgW0=
 nhooyr.io/websocket v1.8.7 h1:usjR2uOr/zjjkVMy0lW+PPohFok7PCow5sDjLgX4P4g=
 nhooyr.io/websocket v1.8.7/go.mod h1:B70DZP8IakI65RVQ51MsWP/8jndNma26DVA/nFSCgW0=
-pgregory.net/rapid v0.4.7 h1:MTNRktPuv5FNqOO151TM9mDTa+XHcX6ypYeISDVD14g=
 pgregory.net/rapid v0.4.7/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU=
 pgregory.net/rapid v1.1.0 h1:CMa0sjHSru3puNx+J0MIAuiiEV4N0qj8/cMWGBBCsjw=
 pgregory.net/rapid v1.1.0/go.mod h1:PY5XlDGj0+V1FCq0o192FdRhpKHGTRIWBgqjDBTrq04=

insecure/go.sum

@@ -2399,8 +2399,8 @@ lukechampine.com/blake3 v1.2.1/go.mod h1:0OFRp7fBtAylGVCO40o87sbupkyIGgbpv1+M1k1
 nhooyr.io/websocket v1.8.6/go.mod h1:B70DZP8IakI65RVQ51MsWP/8jndNma26DVA/nFSCgW0=
 nhooyr.io/websocket v1.8.7 h1:usjR2uOr/zjjkVMy0lW+PPohFok7PCow5sDjLgX4P4g=
 nhooyr.io/websocket v1.8.7/go.mod h1:B70DZP8IakI65RVQ51MsWP/8jndNma26DVA/nFSCgW0=
-pgregory.net/rapid v0.4.7 h1:MTNRktPuv5FNqOO151TM9mDTa+XHcX6ypYeISDVD14g=
 pgregory.net/rapid v0.4.7/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU=
+pgregory.net/rapid v1.1.0 h1:CMa0sjHSru3puNx+J0MIAuiiEV4N0qj8/cMWGBBCsjw=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=

integration/go.sum

@@ -2579,8 +2579,8 @@ modernc.org/memory v1.5.0 h1:N+/8c5rE6EqugZwHii4IFsaJ7MUhoWX07J5tC/iI5Ds=
 modernc.org/memory v1.5.0/go.mod h1:PkUhL0Mugw21sHPeskwZW4D6VscE/GQJOnIpCnW6pSU=
 modernc.org/sqlite v1.21.1 h1:GyDFqNnESLOhwwDRaHGdp2jKLDzpyT/rNLglX3ZkMSU=
 modernc.org/sqlite v1.21.1/go.mod h1:XwQ0wZPIh1iKb5mkvCJ3szzbhk+tykC8ZWqTRTgYRwI=
-pgregory.net/rapid v0.4.7 h1:MTNRktPuv5FNqOO151TM9mDTa+XHcX6ypYeISDVD14g=
 pgregory.net/rapid v0.4.7/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU=
+pgregory.net/rapid v1.1.0 h1:CMa0sjHSru3puNx+J0MIAuiiEV4N0qj8/cMWGBBCsjw=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=
 rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4=
 rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=

module/chainsync/core_rapid_test.go

@@ -25,8 +25,8 @@ func populatedBlockStore(t *rapid.T) []*flow.Header {
 	store := []*flow.Header{unittest.BlockHeaderFixture()}
 	for i := 1; i < NUM_BLOCKS; i++ {
 		// we sample from the store 2/3 times to get deeper trees
-		b := rapid.OneOf(rapid.Just(unittest.BlockHeaderFixture()), rapid.SampledFrom(store), rapid.SampledFrom(store)).Draw(t, "parent").(flow.Header)
-		store = append(store, unittest.BlockHeaderWithParentFixture(&b))
+		b := rapid.OneOf(rapid.Just(unittest.BlockHeaderFixture()), rapid.SampledFrom(store), rapid.SampledFrom(store)).Draw(t, "parent")
+		store = append(store, unittest.BlockHeaderWithParentFixture(b))
 	}
 	return store
 }
@@ -38,8 +38,8 @@ type rapidSync struct {
 	heightRequests map[uint64]bool          // depth 1 pushdown automaton to track height requests
 }
 
-// Init is an action for initializing a rapidSync instance.
-func (r *rapidSync) Init(t *rapid.T) {
+// init is an action for initializing a rapidSync instance.
+func (r *rapidSync) init(t *rapid.T) {
 	var err error
 
 	r.core, err = New(zerolog.New(io.Discard), DefaultConfig(), metrics.NewNoopCollector(), flow.Localnet)
@@ -52,7 +52,7 @@ func (r *rapidSync) Init(t *rapid.T) {
 
 // RequestByID is an action that requests a block by its ID.
 func (r *rapidSync) RequestByID(t *rapid.T) {
-	b := rapid.SampledFrom(r.store).Draw(t, "id_request").(*flow.Header)
+	b := rapid.SampledFrom(r.store).Draw(t, "id_request")
 	r.core.RequestBlock(b.ID(), b.Height)
 	// Re-queueing by ID should always succeed
 	r.idRequests[b.ID()] = true
@@ -62,7 +62,7 @@ func (r *rapidSync) RequestByID(t *rapid.T) {
 
 // RequestByHeight is an action that requests a specific height
 func (r *rapidSync) RequestByHeight(t *rapid.T) {
-	b := rapid.SampledFrom(r.store).Draw(t, "height_request").(*flow.Header)
+	b := rapid.SampledFrom(r.store).Draw(t, "height_request")
 	r.core.RequestHeight(b.Height)
 	// Re-queueing by height should always succeed
 	r.heightRequests[b.Height] = true
@@ -71,8 +71,8 @@ func (r *rapidSync) RequestByHeight(t *rapid.T) {
 // HandleHeight is an action that requests a heights
 // upon receiving an argument beyond a certain tolerance
 func (r *rapidSync) HandleHeight(t *rapid.T) {
-	b := rapid.SampledFrom(r.store).Draw(t, "height_hint_request").(*flow.Header)
-	incr := rapid.IntRange(0, (int)(DefaultConfig().Tolerance)+1).Draw(t, "height increment").(int)
+	b := rapid.SampledFrom(r.store).Draw(t, "height_hint_request")
+	incr := rapid.IntRange(0, (int)(DefaultConfig().Tolerance)+1).Draw(t, "height increment")
 	requestHeight := b.Height + (uint64)(incr)
 	r.core.HandleHeight(b, requestHeight)
 	// Re-queueing by height should always succeed if beyond tolerance
@@ -85,7 +85,7 @@ func (r *rapidSync) HandleHeight(t *rapid.T) {
 
 // HandleByID is an action that provides a block header to the sync engine
 func (r *rapidSync) HandleByID(t *rapid.T) {
-	b := rapid.SampledFrom(r.store).Draw(t, "id_handling").(*flow.Header)
+	b := rapid.SampledFrom(r.store).Draw(t, "id_handling")
 	success := r.core.HandleBlock(b)
 	assert.True(t, success || r.idRequests[b.ID()] == false)
 
@@ -174,7 +174,11 @@ func (r *rapidSync) Check(t *rapid.T) {
 func TestRapidSync(t *testing.T) {
 	unittest.SkipUnless(t, unittest.TEST_FLAKY, "flaky test")
 
-	rapid.Check(t, rapid.Run(&rapidSync{}))
+	rapid.Check(t, func(t *rapid.T) {
+		sm := new(rapidSync)
+		sm.init(t)
+		t.Repeat(rapid.StateMachineActions(sm))
+	})
 }
 
 // utility functions

module/component/component_manager_test.go

@@ -345,7 +345,7 @@ func StartStateTransition() (func(t func()), func(*rapid.T)) {
 
 	executeTransitions := func(t *rapid.T) {
 		for i := 0; i < len(transitions); i++ {
-			j := rapid.IntRange(0, len(transitions)-i-1).Draw(t, "").(int)
+			j := rapid.IntRange(0, len(transitions)-i-1).Draw(t, "")
 			transitions[i], transitions[j+i] = transitions[j+i], transitions[i]
 			transitions[i]()
 		}
@@ -390,35 +390,34 @@ type ComponentManagerMachine struct {
 	assertErrorThrownMatches func(t *rapid.T, err error, msgAndArgs ...interface{})
 	assertErrorNotThrown     func(t *rapid.T)
 
-	cancelGenerator     *rapid.Generator
+	cancelGenerator     *rapid.Generator[bool]
 	drawStateTransition func(t *rapid.T) *StateTransition
 }
 
-func (c *ComponentManagerMachine) Init(t *rapid.T) {
-	numWorkers := rapid.IntRange(0, 5).Draw(t, "num_workers").(int)
-	pCancel := rapid.Float64Range(0, 100).Draw(t, "p_cancel").(float64)
+func (c *ComponentManagerMachine) init(t *rapid.T) {
+	numWorkers := rapid.IntRange(0, 5).Draw(t, "num_workers")
+	pCancel := rapid.Float64Range(0, 100).Draw(t, "p_cancel")
 
-	c.cancelGenerator = rapid.Float64Range(0, 100).
-		Map(func(n float64) bool {
-			return pCancel == 100 || n < pCancel
-		})
+	c.cancelGenerator = rapid.Map(rapid.Float64Range(0, 100), func(n float64) bool {
+		return pCancel == 100 || n < pCancel
+	})
 
 	c.drawStateTransition = func(t *rapid.T) *StateTransition {
 		st := &StateTransition{}
 
 		if !c.canceled {
-			st.cancel = c.cancelGenerator.Draw(t, "cancel").(bool)
+			st.cancel = c.cancelGenerator.Draw(t, "cancel")
 		}
 
 		for workerId, state := range c.workerStates {
 			if allowedTransitions, ok := WorkerStateTransitions[state]; ok {
 				label := fmt.Sprintf("worker_transition_%v", workerId)
 				st.workerIDs = append(st.workerIDs, workerId)
-				st.workerTransitions = append(st.workerTransitions, rapid.SampledFrom(allowedTransitions).Draw(t, label).(WorkerStateTransition))
+				st.workerTransitions = append(st.workerTransitions, rapid.SampledFrom(allowedTransitions).Draw(t, label))
 			}
 		}
 
-		return rapid.Just(st).Draw(t, "state_transition").(*StateTransition)
+		return rapid.Just(st).Draw(t, "state_transition")
 	}
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -625,7 +624,11 @@ func (c *ComponentManagerMachine) Check(t *rapid.T) {
 func TestComponentManager(t *testing.T) {
 	unittest.SkipUnless(t, unittest.TEST_LONG_RUNNING, "skip because this test takes too long")
 
-	rapid.Check(t, rapid.Run(&ComponentManagerMachine{}))
+	rapid.Check(t, func(t *rapid.T) {
+		sm := new(ComponentManagerMachine)
+		sm.init(t)
+		t.Repeat(rapid.StateMachineActions(sm))
+	})
 }
 
 func TestComponentManagerShutdown(t *testing.T) {

module/mempool/stdmap/incorporated_result_seals_test.go

@@ -18,14 +18,14 @@ type icrSealsMachine struct {
 	state []*flow.IncorporatedResultSeal // model of the icrSeals
 }
 
-// Init is an action for initializing a icrSeals instance.
-func (m *icrSealsMachine) Init(t *rapid.T) {
+// init is an action for initializing a icrSeals instance.
+func (m *icrSealsMachine) init(t *rapid.T) {
 	m.icrs = NewIncorporatedResultSeals(1000)
 }
 
 // Add is a conditional action which adds an item to the icrSeals.
 func (m *icrSealsMachine) Add(t *rapid.T) {
-	i := rapid.Uint64().Draw(t, "i").(uint64)
+	i := rapid.Uint64().Draw(t, "i")
 
 	seal := unittest.IncorporatedResultSeal.Fixture(func(s *flow.IncorporatedResultSeal) {
 		s.Header.Height = i
@@ -49,7 +49,7 @@ func (m *icrSealsMachine) Add(t *rapid.T) {
 
 // Prune is a Conditional action that removes elements of height strictly lower than its argument
 func (m *icrSealsMachine) PruneUpToHeight(t *rapid.T) {
-	h := rapid.Uint64().Draw(t, "h").(uint64)
+	h := rapid.Uint64().Draw(t, "h")
 	err := m.icrs.PruneUpToHeight(h)
 	if h >= m.icrs.lowestHeight {
 		require.NoError(t, err)
@@ -72,7 +72,7 @@ func (m *icrSealsMachine) Get(t *rapid.T) {
 	if n == 0 {
 		return
 	}
-	i := rapid.IntRange(0, n-1).Draw(t, "i").(int)
+	i := rapid.IntRange(0, n-1).Draw(t, "i")
 
 	s := m.state[i]
 	actual, ok := m.icrs.ByID(s.ID())
@@ -89,7 +89,7 @@ func (m *icrSealsMachine) GetUnknown(t *rapid.T) {
 	if n == 0 {
 		return
 	}
-	i := rapid.IntRange(0, n-1).Draw(t, "i").(int)
+	i := rapid.IntRange(0, n-1).Draw(t, "i")
 	seal := unittest.IncorporatedResultSeal.Fixture(func(s *flow.IncorporatedResultSeal) {
 		s.Header.Height = uint64(i)
 	})
@@ -117,7 +117,7 @@ func (m *icrSealsMachine) Remove(t *rapid.T) {
 	if n == 0 {
 		return
 	}
-	i := rapid.IntRange(0, n-1).Draw(t, "i").(int)
+	i := rapid.IntRange(0, n-1).Draw(t, "i")
 
 	s := m.state[i]
 	ok := m.icrs.Remove(s.ID())
@@ -137,7 +137,7 @@ func (m *icrSealsMachine) RemoveUnknown(t *rapid.T) {
 	if n == 0 {
 		return
 	}
-	i := rapid.IntRange(0, n-1).Draw(t, "i").(int)
+	i := rapid.IntRange(0, n-1).Draw(t, "i")
 	seal := unittest.IncorporatedResultSeal.Fixture(func(s *flow.IncorporatedResultSeal) {
 		s.Header.Height = uint64(i)
 	})
@@ -168,7 +168,11 @@ func (m *icrSealsMachine) Check(t *rapid.T) {
 
 // Run the icrSeals state machine and test it against its model
 func TestIcrs(t *testing.T) {
-	rapid.Check(t, rapid.Run(&icrSealsMachine{}))
+	rapid.Check(t, func(t *rapid.T) {
+		sm := new(icrSealsMachine)
+		sm.init(t)
+		t.Repeat(rapid.StateMachineActions(sm))
+	})
 }
 
 func TestIncorporatedResultSeals(t *testing.T) {