From b334279635cc8f334af53f1c62cffca31af2216b Mon Sep 17 00:00:00 2001
From: Arran Schlosberg <me@arranschlosberg.com>
Date: Wed, 11 Sep 2024 11:37:39 +0100
Subject: [PATCH] refactor: abstract `hookstest.Register()` out of
 `hookstest.Stub.Register()`

---
 core/state_transition.libevm_test.go |  2 +-
 core/vm/contracts.libevm_test.go     |  6 +++---
 libevm/hookstest/stub.go             | 25 ++++++++++++++++---------
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/core/state_transition.libevm_test.go b/core/state_transition.libevm_test.go
index d81e9ce4f763..fff4267bcf06 100644
--- a/core/state_transition.libevm_test.go
+++ b/core/state_transition.libevm_test.go
@@ -25,7 +25,7 @@ func TestCanExecuteTransaction(t *testing.T) {
 			return makeErr(from, to, s.GetState(account, slot))
 		},
 	}
-	hooks.RegisterForRules(t)
+	hooks.Register(t)
 
 	value := rng.Hash()
 
diff --git a/core/vm/contracts.libevm_test.go b/core/vm/contracts.libevm_test.go
index 28bdcde09a52..585cce03f217 100644
--- a/core/vm/contracts.libevm_test.go
+++ b/core/vm/contracts.libevm_test.go
@@ -64,7 +64,7 @@ func TestPrecompileOverride(t *testing.T) {
 					},
 				},
 			}
-			hooks.RegisterForRules(t)
+			hooks.Register(t)
 
 			t.Run(fmt.Sprintf("%T.Call([overridden precompile address = %v])", &vm.EVM{}, tt.addr), func(t *testing.T) {
 				_, evm := ethtest.NewZeroEVM(t)
@@ -103,7 +103,7 @@ func TestNewStatefulPrecompile(t *testing.T) {
 			),
 		},
 	}
-	hooks.RegisterForRules(t)
+	hooks.Register(t)
 
 	caller := rng.Address()
 	input := rng.Bytes(8)
@@ -133,7 +133,7 @@ func TestCanCreateContract(t *testing.T) {
 			return makeErr(cc, s.GetState(account, slot))
 		},
 	}
-	hooks.RegisterForRules(t)
+	hooks.Register(t)
 
 	origin := rng.Address()
 	caller := rng.Address()
diff --git a/libevm/hookstest/stub.go b/libevm/hookstest/stub.go
index 3dcb1619c3cd..e8fda4310781 100644
--- a/libevm/hookstest/stub.go
+++ b/libevm/hookstest/stub.go
@@ -1,4 +1,5 @@
-// Package hookstest provides test doubles for testing subsets of libevm hooks.
+// Package hookstest provides test doubles and convenience wrappers for testing
+// libevm hooks.
 package hookstest
 
 import (
@@ -10,6 +11,15 @@ import (
 	"github.com/ethereum/go-ethereum/params"
 )
 
+// Register clears any registered [params.Extras] and then registers `extras`
+// for the liftime of the current test, clearing them via tb's
+// [testing.TB.Cleanup].
+func Register[C params.ChainConfigHooks, R params.RulesHooks](tb testing.TB, extras params.Extras[C, R]) {
+	params.TestOnlyClearRegisteredExtras()
+	tb.Cleanup(params.TestOnlyClearRegisteredExtras)
+	params.RegisterExtras(extras)
+}
+
 // A Stub is a test double for [params.ChainConfigHooks] and
 // [params.RulesHooks]. Each of the fields, if non-nil, back their respective
 // hook methods, which otherwise fall back to the default behaviour.
@@ -19,17 +29,14 @@ type Stub struct {
 	CanCreateContractFn     func(*libevm.AddressContext, libevm.StateReader) error
 }
 
-// RegisterForRules clears any registered [params.Extras] and then registers s
-// as [params.RulesHooks], which are themselves cleared by the
-// [testing.TB.Cleanup] routine.
-func (s *Stub) RegisterForRules(tb testing.TB) {
-	params.TestOnlyClearRegisteredExtras()
-	params.RegisterExtras(params.Extras[params.NOOPHooks, Stub]{
-		NewRules: func(_ *params.ChainConfig, _ *params.Rules, _ *params.NOOPHooks, blockNum *big.Int, isMerge bool, timestamp uint64) *Stub {
+// Register is a convenience wrapper for registering s as both the
+// [params.ChainConfigHooks] and [params.RulesHooks] via [Register].
+func (s *Stub) Register(tb testing.TB) {
+	Register(tb, params.Extras[Stub, Stub]{
+		NewRules: func(_ *params.ChainConfig, _ *params.Rules, _ *Stub, blockNum *big.Int, isMerge bool, timestamp uint64) *Stub {
 			return s
 		},
 	})
-	tb.Cleanup(params.TestOnlyClearRegisteredExtras)
 }
 
 func (s Stub) PrecompileOverride(a common.Address) (libevm.PrecompiledContract, bool) {