diff --git a/mock/contracts/deployerSimpleSC.go b/mock/contracts/deployerSimpleSC.go
index 508df552e..8f73cc361 100644
--- a/mock/contracts/deployerSimpleSC.go
+++ b/mock/contracts/deployerSimpleSC.go
@@ -57,6 +57,12 @@ func InitMockMethod(instanceMock *mock.InstanceMock, config interface{}) {
 	instanceMock.AddMockMethod("init", testcommon.SimpleWasteGasMockMethod(instanceMock, testConfig.GasUsedByInit))
 }
 
+// UpgradeMockMethod -
+func UpgradeMockMethod(instanceMock *mock.InstanceMock, config interface{}) {
+	testConfig := config.(*testcommon.TestConfig)
+	instanceMock.AddMockMethod("upgrade", testcommon.SimpleWasteGasMockMethod(instanceMock, testConfig.GasUsedByInit))
+}
+
 // CallbackTestConfig -
 type CallbackTestConfig interface {
 	CallbackFails() bool
diff --git a/test/contracts/init-correct/init-correct.c b/test/contracts/init-correct/init-correct.c
index cce4f2a8d..6f2023c62 100644
--- a/test/contracts/init-correct/init-correct.c
+++ b/test/contracts/init-correct/init-correct.c
@@ -27,3 +27,7 @@ void init() {
 		}
 	}
 }
+
+void upgrade() {
+	init();
+}
diff --git a/test/contracts/init-correct/init-correct.export b/test/contracts/init-correct/init-correct.export
index b1b716105..e4e16e74c 100644
--- a/test/contracts/init-correct/init-correct.export
+++ b/test/contracts/init-correct/init-correct.export
@@ -1 +1,2 @@
 init
+upgrade
diff --git a/test/contracts/init-correct/output/init-correct.wasm b/test/contracts/init-correct/output/init-correct.wasm
index eb66323d6..afd4198d4 100755
Binary files a/test/contracts/init-correct/output/init-correct.wasm and b/test/contracts/init-correct/output/init-correct.wasm differ
diff --git a/vmhost/hosttest/contracts_deploy_test.go b/vmhost/hosttest/contracts_deploy_test.go
index 8b6c946d7..792d44dbb 100644
--- a/vmhost/hosttest/contracts_deploy_test.go
+++ b/vmhost/hosttest/contracts_deploy_test.go
@@ -386,9 +386,14 @@ func runUpdateFromSourceTest(t *testing.T, testConfig *testcommon.TestConfig, as
 	var deployedContract test.MockTestSmartContract
 	var contractToUpdate test.MockTestSmartContract
 	if testConfig.DeployedContractAddress != nil {
+		// The DeployedContract will be the source of the code to be written in
+		// ContractToBeUpdated. Therefore it is the DeployedContract which must
+		// initially have UpgradeMockMethod. After the code update, the
+		// ContractToBeUpdated will also have the UpgradeMockMethod and it will be
+		// called by the VM.
 		deployedContract = test.CreateMockContract(testConfig.DeployedContractAddress).
 			WithConfig(testConfig).
-			WithMethods(contracts.InitMockMethod)
+			WithMethods(contracts.InitMockMethod, contracts.UpgradeMockMethod)
 	}
 	if testConfig.ContractToBeUpdatedAddress != nil {
 		contractToUpdate = test.CreateMockContract(testConfig.ContractToBeUpdatedAddress).
diff --git a/vmhost/hosttest/execution_test.go b/vmhost/hosttest/execution_test.go
index fc74f012f..3458ca126 100644
--- a/vmhost/hosttest/execution_test.go
+++ b/vmhost/hosttest/execution_test.go
@@ -185,7 +185,7 @@ func TestExecution_DeployWASM_Successful(t *testing.T) {
 		AndAssertResults(func(blockchainHook *contextmock.BlockchainHookStub, verify *test.VMOutputVerifier) {
 			verify.Ok().
 				ReturnData([]byte("init successful")).
-				GasRemaining(430).
+				GasRemaining(410).
 				Nonce([]byte("caller"), 24).
 				Code(newAddress, input.ContractCode).
 				BalanceDelta(newAddress, 88)
@@ -449,7 +449,7 @@ func TestExecution_Deploy_DisallowFloatingPoint(t *testing.T) {
 
 func TestExecution_DeployWASM_GasValidation(t *testing.T) {
 	var gasProvided uint64
-	gasUsedByDeployment := uint64(570)
+	gasUsedByDeployment := uint64(590)
 
 	inputBuilder := test.CreateTestContractCreateInputBuilder().
 		WithContractCode(test.GetTestSCCode("init-correct", "../../")).
@@ -2983,13 +2983,13 @@ func TestExecution_CreateNewContract_Success(t *testing.T) {
 		AndAssertResults(func(host vmhost.VMHost, stubBlockchainHook *contextmock.BlockchainHookStub, verify *test.VMOutputVerifier) {
 			verify.Ok().
 				Balance(test.ParentAddress, 1000).
-				GasUsed(test.ParentAddress, 1069).
-				GasRemaining(998361).
+				GasUsed(test.ParentAddress, 1109).
+				GasRemaining(998301).
 				BalanceDelta(childAddress, 42).
 				Code(childAddress, childCode).
 				CodeMetadata(childAddress, []byte{1, 0}).
 				CodeDeployerAddress(childAddress, test.ParentAddress).
-				GasUsed(childAddress, 570).
+				GasUsed(childAddress, 590).
 				ReturnData([]byte{byte(l / 256), byte(l % 256)}, []byte("init successful"), []byte("succ")).
 				Storage().
 				Logs(vmcommon.LogEntry{