From 9d0a4522d4b33356ddbb59fc16e22c7d619b9d84 Mon Sep 17 00:00:00 2001
From: Damian Nolan <damiannolan@gmail.com>
Date: Sun, 11 Jun 2023 13:20:23 +0200
Subject: [PATCH] adding testing app overrides for upgrade handshake app
 callback handlers

---
 testing/mock/ibc_app.go    | 36 ++++++++++++++++++++++++++++++++++++
 testing/mock/ibc_module.go | 20 ++++++++++++++++++++
 testing/simapp/app.go      |  2 ++
 3 files changed, 58 insertions(+)

diff --git a/testing/mock/ibc_app.go b/testing/mock/ibc_app.go
index d3bac8889a2..e7c47153608 100644
--- a/testing/mock/ibc_app.go
+++ b/testing/mock/ibc_app.go
@@ -85,6 +85,42 @@ type IBCApp struct {
 		packet channeltypes.Packet,
 		relayer sdk.AccAddress,
 	) error
+
+	OnChanUpgradeInit func(
+		ctx sdk.Context,
+		portID, channelID string,
+		order channeltypes.Order,
+		connectionHops []string,
+		sequence uint64,
+		version, previousVersion string,
+	) (string, error)
+
+	OnChanUpgradeTry func(
+		ctx sdk.Context,
+		portID, channelID string,
+		order channeltypes.Order,
+		connectionHops []string,
+		counterpartyVersion string,
+	) (string, error)
+
+	OnChanUpgradeAck func(
+		ctx sdk.Context,
+		portID,
+		channelID,
+		counterpartyVersion string,
+	) error
+
+	OnChanUpgradeOpen func(
+		ctx sdk.Context,
+		portID,
+		channelID string,
+	) error
+
+	OnChanUpgradeRestore func(
+		ctx sdk.Context,
+		portID,
+		channelID string,
+	) error
 }
 
 // NewIBCApp returns a IBCApp. An empty PortID indicates the mock app doesn't bind/claim ports.
diff --git a/testing/mock/ibc_module.go b/testing/mock/ibc_module.go
index 49385afc415..59e12c03e90 100644
--- a/testing/mock/ibc_module.go
+++ b/testing/mock/ibc_module.go
@@ -164,26 +164,46 @@ func (im IBCModule) OnTimeoutPacket(ctx sdk.Context, packet channeltypes.Packet,
 
 // OnChanUpgradeInit implements the IBCModule interface
 func (im IBCModule) OnChanUpgradeInit(ctx sdk.Context, portID, channelID string, order channeltypes.Order, connectionHops []string, sequence uint64, version, previousVersion string) (string, error) {
+	if im.IBCApp.OnChanUpgradeInit != nil {
+		return im.IBCApp.OnChanUpgradeInit(ctx, portID, channelID, order, connectionHops, sequence, version, previousVersion)
+	}
+
 	return version, nil
 }
 
 // OnChanUpgradeTry implements the IBCModule interface
 func (im IBCModule) OnChanUpgradeTry(ctx sdk.Context, portID, channelID string, order channeltypes.Order, connectionHops []string, counterpartyVersion string) (string, error) {
+	if im.IBCApp.OnChanUpgradeTry != nil {
+		return im.IBCApp.OnChanUpgradeTry(ctx, portID, channelID, order, connectionHops, counterpartyVersion)
+	}
+
 	return counterpartyVersion, nil
 }
 
 // OnChanUpgradeAck implements the IBCModule interface
 func (im IBCModule) OnChanUpgradeAck(ctx sdk.Context, portID, channelID, counterpartyVersion string) error {
+	if im.IBCApp.OnChanUpgradeAck != nil {
+		return im.IBCApp.OnChanUpgradeAck(ctx, portID, channelID, counterpartyVersion)
+	}
+
 	return nil
 }
 
 // OnChanUpgradeOpen implements the IBCModule interface
 func (im IBCModule) OnChanUpgradeOpen(ctx sdk.Context, portID, channelID string) error {
+	if im.IBCApp.OnChanUpgradeOpen != nil {
+		return im.IBCApp.OnChanUpgradeOpen(ctx, portID, channelID)
+	}
+
 	return nil
 }
 
 // OnChanUpgradeRestore implements the IBCModule interface
 func (im IBCModule) OnChanUpgradeRestore(ctx sdk.Context, portID, channelID string) error {
+	if im.IBCApp.OnChanUpgradeRestore != nil {
+		return im.IBCApp.OnChanUpgradeRestore(ctx, portID, channelID)
+	}
+
 	return nil
 }
 
diff --git a/testing/simapp/app.go b/testing/simapp/app.go
index 6e97d76c2fc..ac3654e4117 100644
--- a/testing/simapp/app.go
+++ b/testing/simapp/app.go
@@ -245,6 +245,7 @@ type SimApp struct {
 
 	// make IBC modules public for test purposes
 	// these modules are never directly routed to by the IBC Router
+	IBCMockModule ibcmock.IBCModule
 	ICAAuthModule ibcmock.IBCModule
 	FeeMockModule ibcmock.IBCModule
 
@@ -455,6 +456,7 @@ func NewSimApp(
 
 	// The mock module is used for testing IBC
 	mockIBCModule := ibcmock.NewIBCModule(&mockModule, ibcmock.NewIBCApp(ibcmock.ModuleName, scopedIBCMockKeeper))
+	app.IBCMockModule = mockIBCModule
 	ibcRouter.AddRoute(ibcmock.ModuleName, mockIBCModule)
 
 	// Create Transfer Stack