Add decorator options

x/wasm/keeper/options.go

@@ -21,23 +21,39 @@ func WithWasmEngine(x types.WasmerEngine) Option {
 }
 
 // WithMessageHandler is an optional constructor parameter to set a custom handler for wasmVM messages.
-// This option should not be combined with Option `WithMessageEncoders`.
+// This option should not be combined with Option `WithMessageEncoders` or `WithMessageHandlerDecorator`
 func WithMessageHandler(x Messenger) Option {
 	return optsFn(func(k *Keeper) {
 		k.messenger = x
 	})
 }
 
+// WithMessageHandlerDecorator is an optional constructor parameter to decorate the wasm handler for wasmVM messages.
+// This option should not be combined with Option `WithMessageEncoders` or `WithMessageHandler`
+func WithMessageHandlerDecorator(d func(old Messenger) Messenger) Option {
+	return optsFn(func(k *Keeper) {
+		k.messenger = d(k.messenger)
+	})
+}
+
 // WithQueryHandler is an optional constructor parameter to set custom query handler for wasmVM requests.
-// This option should not be combined with Option `WithQueryPlugins`.
+// This option should not be combined with Option `WithQueryPlugins` or `WithQueryHandlerDecorator`
 func WithQueryHandler(x WasmVMQueryHandler) Option {
 	return optsFn(func(k *Keeper) {
 		k.wasmVMQueryHandler = x
 	})
 }
 
+// WithQueryHandlerDecorator is an optional constructor parameter to decorate the default wasm query handler for wasmVM requests.
+// This option should not be combined with Option `WithQueryPlugins` or `WithQueryHandler`
+func WithQueryHandlerDecorator(d func(old WasmVMQueryHandler) WasmVMQueryHandler) Option {
+	return optsFn(func(k *Keeper) {
+		k.wasmVMQueryHandler = d(k.wasmVMQueryHandler)
+	})
+}
+
 // WithQueryPlugins is an optional constructor parameter to pass custom query plugins for wasmVM requests.
-// This option expects the default `QueryHandler` set an should not be combined with Option `WithQueryHandler`.
+// This option expects the default `QueryHandler` set an should not be combined with Option `WithQueryHandler` or `WithQueryHandlerDecorator`.
 func WithQueryPlugins(x *QueryPlugins) Option {
 	return optsFn(func(k *Keeper) {
 		q, ok := k.wasmVMQueryHandler.(QueryPlugins)
@@ -49,7 +65,7 @@ func WithQueryPlugins(x *QueryPlugins) Option {
 }
 
 // WithMessageEncoders is an optional constructor parameter to pass custom message encoder to the default wasm message handler.
-// This option expects the `DefaultMessageHandler` set an should not be combined with Option `WithMessageHandler`.
+// This option expects the `DefaultMessageHandler` set and should not be combined with Option `WithMessageHandler` or `WithMessageHandlerDecorator`.
 func WithMessageEncoders(x *MessageEncoders) Option {
 	return optsFn(func(k *Keeper) {
 		q, ok := k.messenger.(*MessageHandlerChain)

x/wasm/keeper/options_test.go

@@ -8,6 +8,7 @@ import (
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"testing"
 )
 
@@ -19,31 +20,49 @@ func TestConstructorOptions(t *testing.T) {
 		"wasm engine": {
 			srcOpt: WithWasmEngine(&wasmtesting.MockWasmer{}),
 			verify: func(t *testing.T, k Keeper) {
-				assert.IsType(t, k.wasmVM, &wasmtesting.MockWasmer{})
+				assert.IsType(t, &wasmtesting.MockWasmer{}, k.wasmVM)
 			},
 		},
 		"message handler": {
 			srcOpt: WithMessageHandler(&wasmtesting.MockMessageHandler{}),
 			verify: func(t *testing.T, k Keeper) {
-				assert.IsType(t, k.messenger, &wasmtesting.MockMessageHandler{})
+				assert.IsType(t, &wasmtesting.MockMessageHandler{}, k.messenger)
 			},
 		},
 		"query plugins": {
 			srcOpt: WithQueryHandler(&wasmtesting.MockQueryHandler{}),
 			verify: func(t *testing.T, k Keeper) {
-				assert.IsType(t, k.wasmVMQueryHandler, &wasmtesting.MockQueryHandler{})
+				assert.IsType(t, &wasmtesting.MockQueryHandler{}, k.wasmVMQueryHandler)
+			},
+		},
+		"message handler decorator": {
+			srcOpt: WithMessageHandlerDecorator(func(old Messenger) Messenger {
+				require.IsType(t, &MessageHandlerChain{}, old)
+				return &wasmtesting.MockMessageHandler{}
+			}),
+			verify: func(t *testing.T, k Keeper) {
+				assert.IsType(t, &wasmtesting.MockMessageHandler{}, k.messenger)
+			},
+		},
+		"query plugins decorator": {
+			srcOpt: WithQueryHandlerDecorator(func(old WasmVMQueryHandler) WasmVMQueryHandler {
+				require.IsType(t, QueryPlugins{}, old)
+				return &wasmtesting.MockQueryHandler{}
+			}),
+			verify: func(t *testing.T, k Keeper) {
+				assert.IsType(t, &wasmtesting.MockQueryHandler{}, k.wasmVMQueryHandler)
 			},
 		},
 		"coin transferrer": {
 			srcOpt: WithCoinTransferrer(&wasmtesting.MockCoinTransferrer{}),
 			verify: func(t *testing.T, k Keeper) {
-				assert.IsType(t, k.bank, &wasmtesting.MockCoinTransferrer{})
+				assert.IsType(t, &wasmtesting.MockCoinTransferrer{}, k.bank)
 			},
 		},
 		"costs": {
 			srcOpt: WithGasRegister(&wasmtesting.MockGasRegister{}),
 			verify: func(t *testing.T, k Keeper) {
-				assert.IsType(t, k.gasRegister, &wasmtesting.MockGasRegister{})
+				assert.IsType(t, &wasmtesting.MockGasRegister{}, k.gasRegister)
 			},
 		},
 		"api costs": {

x/wasm/keeper/query_plugins.go

@@ -493,3 +493,13 @@ func convertSdkCoinToWasmCoin(coin sdk.Coin) wasmvmtypes.Coin {
 		Amount: coin.Amount.String(),
 	}
 }
+
+var _ WasmVMQueryHandler = WasmVMQueryHandlerFn(nil)
+
+// WasmVMQueryHandlerFn is a helper to construct a function based query handler.
+type WasmVMQueryHandlerFn func(ctx sdk.Context, caller sdk.AccAddress, request wasmvmtypes.QueryRequest) ([]byte, error)
+
+// HandleQuery delegates call into wrapped WasmVMQueryHandlerFn
+func (w WasmVMQueryHandlerFn) HandleQuery(ctx sdk.Context, caller sdk.AccAddress, request wasmvmtypes.QueryRequest) ([]byte, error) {
+	return w(ctx, caller, request)
+}

x/wasm/keeper/recurse_test.go

@@ -2,7 +2,6 @@ package keeper
 
 import (
 	"encoding/json"
-	"github.com/CosmWasm/wasmd/x/wasm/keeper/wasmtesting"
 	"github.com/CosmWasm/wasmd/x/wasm/types"
 	"testing"
 
@@ -39,18 +38,14 @@ type recurseResponse struct {
 var totalWasmQueryCounter int
 
 func initRecurseContract(t *testing.T) (contract sdk.AccAddress, creator sdk.AccAddress, ctx sdk.Context, keeper *Keeper) {
-	// we do one basic setup before all test cases (which are read-only and don't change state)
-	var realWasmQuerier func(ctx sdk.Context, request *wasmvmtypes.WasmQuery) ([]byte, error)
-	countingQuerier := &wasmtesting.MockQueryHandler{
-		HandleQueryFn: func(ctx sdk.Context, request wasmvmtypes.QueryRequest, caller sdk.AccAddress) ([]byte, error) {
+	countingQuerierDec := func(realWasmQuerier WasmVMQueryHandler) WasmVMQueryHandler {
+		return WasmVMQueryHandlerFn(func(ctx sdk.Context, caller sdk.AccAddress, request wasmvmtypes.QueryRequest) ([]byte, error) {
 			totalWasmQueryCounter++
-			return realWasmQuerier(ctx, request.Wasm)
-		}}
-
-	ctx, keepers := CreateTestInput(t, false, SupportedFeatures, WithQueryHandler(countingQuerier))
+			return realWasmQuerier.HandleQuery(ctx, caller, request)
+		})
+	}
+	ctx, keepers := CreateTestInput(t, false, SupportedFeatures, WithQueryHandlerDecorator(countingQuerierDec))
 	keeper = keepers.WasmKeeper
-	realWasmQuerier = WasmQuerier(keeper)
-
 	exampleContract := InstantiateHackatomExampleContract(t, ctx, keepers)
 	return exampleContract.Contract, exampleContract.CreatorAddr, ctx, keeper
 }