[Feature]: Add support for test sdk.Context with n store keys, plus mem store keys

[damiannolan]

Summary

Add support for test sdk.Context to be able to mount any number (and supported type - key, transientKey, memKey) of keys on the cms before creating and returning the test context.

Problem Definition

This is a super nice test helper for creating a sdk.Context with a custom store key - https://github.com/cosmos/cosmos-sdk/blob/main/testutil/context.go#L20.

However, it doesn't support mounting multiple store keys at once and the ability to mount MemoryStoreKeys on the cms.

It's likely no one had faced an issue with this as module's tend to have a single store key. However I ran into an issue when refactoring tests for the sdk v0.50 upgrade, as x/capability deals with mem store keys. See cosmos/ibc-go#4068 (comment)

Proposal

Add something like this as to not break the current usage in existing tests. Or rip the bandaid off and change the DefaultContext API.

func DefaultContextWithKeys(
	keys map[string]*storetypes.KVStoreKey,
	transKeys map[string]*storetypes.TransientStoreKey,
	memKeys map[string]*storetypes.MemoryStoreKey,
) sdk.Context {
	db := dbm.NewMemDB()
	cms := store.NewCommitMultiStore(db, log.NewNopLogger(), metrics.NewNoOpMetrics())

	for _, key := range keys {
		cms.MountStoreWithDB(key, storetypes.StoreTypeIAVL, db)
	}

	for _, tKey := range transKeys {
		cms.MountStoreWithDB(tKey, storetypes.StoreTypeTransient, db)
	}

	for _, memkey := range memKeys {
		cms.MountStoreWithDB(memkey, storetypes.StoreTypeMemory, db)
	}

	err := cms.LoadLatestVersion()
	if err != nil {
		panic(err)
	}

	return sdk.NewContext(cms, cmtproto.Header{}, false, log.NewNopLogger())
}

Usage looks something like:

func TestFoo() {
	keys := storetypes.NewKVStoreKeys("mock")
	tKeys := storetypes.NewTransientStoreKeys("t-mock")
	memKeys := storetypes.NewMemoryStoreKeys("mem-mock", "mem-mock-2")

	ctx := DefaultContextWithKeys(keys, tKeys, memKeys)
}

[julienrbrt]

Seems useful to me, do you want a submit a PR for that?

[damiannolan]

Yeah definitely, I'll try get to it next week when I'm back around