From 382de3311c057ac23acf5689801fb7023bbef1f4 Mon Sep 17 00:00:00 2001
From: Damian Nolan <damiannolan@gmail.com>
Date: Mon, 31 Jul 2023 14:26:47 +0200
Subject: [PATCH] feat(testutil): adding DefaultContextWithKeys test helper
 (#17216)

---
 testutil/context.go | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/testutil/context.go b/testutil/context.go
index 399778ee5a76..876b6a7b7508 100644
--- a/testutil/context.go
+++ b/testutil/context.go
@@ -31,6 +31,36 @@ func DefaultContext(key, tkey storetypes.StoreKey) sdk.Context {
 	return ctx
 }
 
+// DefaultContextWithKeys creates a sdk.Context with a fresh MemDB, mounting the providing keys for usage in the multistore.
+// This function is intended to be used for testing purposes only.
+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())
+}
+
 type TestContext struct {
 	Ctx sdk.Context
 	DB  *dbm.MemDB