Add wasmvm decorator option

CosmWasm · Apr 19, 2023 · e5df7cb · e5df7cb
1 parent 88439cb
commit e5df7cb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/x/wasm/keeper/options.go b/x/wasm/keeper/options.go
@@ -24,6 +24,13 @@ func WithWasmEngine(x types.WasmerEngine) Option {
 	})
 }
 
+// WithWasmEngineDecorator is an optional constructor parameter to decorate the default wasmVM engine.
+func WithWasmEngineDecorator(d func(old types.WasmerEngine) types.WasmerEngine) Option {
+	return optsFn(func(k *Keeper) {
+		k.wasmVM = d(k.wasmVM)
+	})
+}
+
 // WithMessageHandler is an optional constructor parameter to set a custom handler for wasmVM messages.
 // This option should not be combined with Option `WithMessageEncoders` or `WithMessageHandlerDecorator`
 func WithMessageHandler(x Messenger) Option {

diff --git a/x/wasm/keeper/options_test.go b/x/wasm/keeper/options_test.go
@@ -4,6 +4,8 @@ import (
 	"reflect"
 	"testing"
 
+	wasmvm "github.com/CosmWasm/wasmvm"
+
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
@@ -27,6 +29,15 @@ func TestConstructorOptions(t *testing.T) {
 				assert.IsType(t, &wasmtesting.MockWasmer{}, k.wasmVM)
 			},
 		},
+		"decorate wasmvm": {
+			srcOpt: WithWasmEngineDecorator(func(old types.WasmerEngine) types.WasmerEngine {
+				require.IsType(t, &wasmvm.VM{}, old)
+				return &wasmtesting.MockWasmer{}
+			}),
+			verify: func(t *testing.T, k Keeper) {
+				assert.IsType(t, &wasmtesting.MockWasmer{}, k.wasmVM)
+			},
+		},
 		"message handler": {
 			srcOpt: WithMessageHandler(&wasmtesting.MockMessageHandler{}),
 			verify: func(t *testing.T, k Keeper) {