diff --git a/contracts/test/TestContractManager.sol b/contracts/test/TestContractManager.sol
index 036b7b6c3..6595a7b35 100644
--- a/contracts/test/TestContractManager.sol
+++ b/contracts/test/TestContractManager.sol
@@ -38,6 +38,10 @@ contract ContractManager is IContractManagerTester {
 
     event ContractUpgraded(string contractsName, address contractsAddress);
 
+    error ContractNotFound(
+        string contractName
+    );
+
     constructor() {
         owner = msg.sender;
     }
@@ -73,6 +77,8 @@ contract ContractManager is IContractManagerTester {
         returns (address contractAddress)
     {
         contractAddress = contracts[keccak256(abi.encodePacked(name))];
-        require(contractAddress != address(0), "Contract has not been found");
+        if (contractAddress == address(0)) {
+            revert ContractNotFound(name);
+        }
     }
 }
diff --git a/migrations/upgradeMainnet.ts b/migrations/upgradeMainnet.ts
index 4a9714c4d..733c5ed44 100644
--- a/migrations/upgradeMainnet.ts
+++ b/migrations/upgradeMainnet.ts
@@ -159,13 +159,12 @@ async function updateAbi() {
 }
 
 async function main() {
-    let contractNamesToUpgrade = contracts;
-    if (process.env.TEST_UPGRADE !== "true") {
-        contractNamesToUpgrade = [
-            "MessageProxyForMainnet",
-            "CommunityPool"
-        ]
-
+    let contractNamesToUpgrade = [
+        "MessageProxyForMainnet",
+        "CommunityPool"
+    ]
+    if (process.env.UPGRADE_ALL) {
+        contractNamesToUpgrade = contracts;
     }
     const upgrader = new ImaMainnetUpgrader(
         "2.1.0",
diff --git a/scripts/test_upgrade.sh b/scripts/test_upgrade.sh
index 8442ef58c..ada74dd73 100755
--- a/scripts/test_upgrade.sh
+++ b/scripts/test_upgrade.sh
@@ -71,7 +71,7 @@ cd "$GITHUB_WORKSPACE"
 rm -r --interactive=never "$DEPLOYED_DIR"
 
 MESSAGE_PROXY_FOR_MAINNET=$(cat data/proxyMainnet.json | jq -r .message_proxy_mainnet_address)
-TEST_UPGRADE=true \
+UPGRADE_ALL=true \
 ABI="data/proxyMainnet.json" \
 TARGET="$MESSAGE_PROXY_FOR_MAINNET" \
 ALLOW_NOT_ATOMIC_UPGRADE="OK" \