diff --git a/contracts/test/TestContractManager.sol b/contracts/test/TestContractManager.sol
index aca6a39bf..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;
     }
@@ -64,7 +68,17 @@ contract ContractManager is IContractManagerTester {
     /**
      * @dev Returns the contract address for a given contractName.
      */
-    function getContract(string memory contractName) external view override returns (address) {
-        return contracts[keccak256(abi.encodePacked(contractName))];
+    function getContract(
+        string memory name
+    )
+        public
+        view
+        override
+        returns (address contractAddress)
+    {
+        contractAddress = contracts[keccak256(abi.encodePacked(name))];
+        if (contractAddress == address(0)) {
+            revert ContractNotFound(name);
+        }
     }
 }
diff --git a/migrations/upgradeMainnet.ts b/migrations/upgradeMainnet.ts
index 226646d10..733c5ed44 100644
--- a/migrations/upgradeMainnet.ts
+++ b/migrations/upgradeMainnet.ts
@@ -103,19 +103,27 @@ class ImaMainnetUpgrader extends Upgrader {
             contractManagerInterface,
             ethers.provider
         )
+        if (!process.env.ABI) {
+            console.log(chalk.red("Set path to file with ABI and addresses to ABI environment variables"));
+            process.exit(1);
+        }
+        const abi = JSON.parse(await fs.readFile(process.env.ABI, "utf-8"));
         for (const contractName of contracts) {
             try {
                 const contractAddress = await contractManager.getContract(contractName);
                 console.log(`Address of ${contractName} is set to ${contractAddress}`);
             } catch {
                 // getContract failed because the contract is not set
-                const contractAddress = await this.instance.getContract(contractName);
+                const contractAddress = abi[`${getContractKeyInAbiFile(contractName)}_address`] as string;
                 this.transactions.push(Transaction.from(
                     {
                         to: await contractManager.getAddress(),
                         data: contractManager.interface.encodeFunctionData(
                             "setContractsAddress",
-                            [contractAddress]
+                            [
+                                contractName,
+                                contractAddress
+                            ]
                         )
                     }
                 ))
@@ -151,10 +159,17 @@ async function updateAbi() {
 }
 
 async function main() {
+    let contractNamesToUpgrade = [
+        "MessageProxyForMainnet",
+        "CommunityPool"
+    ]
+    if (process.env.UPGRADE_ALL) {
+        contractNamesToUpgrade = contracts;
+    }
     const upgrader = new ImaMainnetUpgrader(
         "2.1.0",
         await getImaMainnetInstance(),
-        contracts
+        contractNamesToUpgrade
     );
     await upgrader.upgrade();
     await updateAbi();
diff --git a/migrations/upgradeSchain.ts b/migrations/upgradeSchain.ts
index 57f8ea942..fd7076722 100644
--- a/migrations/upgradeSchain.ts
+++ b/migrations/upgradeSchain.ts
@@ -92,10 +92,15 @@ async function updateAbi(contracts: string[]) {
 async function main() {
     const pathToManifest: string = process.env.MANIFEST || "";
     await manifestSetup(pathToManifest);
+    let contractNamesToUpgrade: string[] = [
+    ]
+    if (process.env.UPGRADE_ALL) {
+        contractNamesToUpgrade = contracts;
+    }
     const upgrader = new ImaSchainUpgrader(
         "2.1.0",
         await getImaSchainInstance(),
-        contracts
+        contractNamesToUpgrade
     );
     await upgrader.upgrade();
     updateAbi(contracts);
diff --git a/scripts/test_upgrade.sh b/scripts/test_upgrade.sh
index 8442ef58c..93201338f 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" \
@@ -87,6 +87,7 @@ MESSAGE_PROXY_FOR_SCHAIN=$(cat data/$ABI_FILENAME_SCHAIN | jq -r .message_proxy_
 # ABI="data/$ABI_FILENAME_SCHAIN" \
 # MANIFEST="data/ima-schain-$DEPLOYED_VERSION-manifest.json" \
 # CHAIN_NAME_SCHAIN="Test" \
+# UPGRADE_ALL=true \
 # ALLOW_NOT_ATOMIC_UPGRADE="OK" \
 # TARGET="$MESSAGE_PROXY_FOR_SCHAIN" \
 # VERSION=$VERSION_TAG \
diff --git a/test/utils/deploy/mainnet/messageProxyForMainnet.ts b/test/utils/deploy/mainnet/messageProxyForMainnet.ts
index e169c9312..a14a6b6a3 100644
--- a/test/utils/deploy/mainnet/messageProxyForMainnet.ts
+++ b/test/utils/deploy/mainnet/messageProxyForMainnet.ts
@@ -7,11 +7,12 @@ export async function deployMessageProxyForMainnet(
     contractManager: ContractManager
 ) {
     const factory = await ethers.getContractFactory(name);
-    if (await contractManager.getContract(name) !== "0x0000000000000000000000000000000000000000") {
-        return factory.attach(await contractManager.getContract(name)) as MessageProxyForMainnet;
-    } else {
+    try {
+        await contractManager.getContract(name);
+    } catch {
         const instance = await upgrades.deployProxy(factory, [await contractManager.getAddress()]) as unknown as MessageProxyForMainnet;
         await contractManager.setContractsAddress(name, instance);
         return instance;
     }
+    return factory.attach(await contractManager.getContract(name)) as MessageProxyForMainnet;
 }
diff --git a/test/utils/deploy/test/messageProxyForMainnetTester.ts b/test/utils/deploy/test/messageProxyForMainnetTester.ts
index 2a2ac35eb..0053f02e1 100644
--- a/test/utils/deploy/test/messageProxyForMainnetTester.ts
+++ b/test/utils/deploy/test/messageProxyForMainnetTester.ts
@@ -7,9 +7,9 @@ export async function deployMessageProxyForMainnetTester(
     contractManager: ContractManager
 ) {
     const factory = await ethers.getContractFactory(name);
-    if (await contractManager.getContract(name) !== "0x0000000000000000000000000000000000000000") {
-        return factory.attach(await contractManager.getContract(name)) as MessageProxyForMainnetTester;
-    } else {
+    try {
+        await contractManager.getContract(name);
+    } catch {
         const instance = await upgrades.deployProxy(
             factory,
             [await contractManager.getAddress()]
@@ -17,4 +17,5 @@ export async function deployMessageProxyForMainnetTester(
         await contractManager.setContractsAddress(name, instance);
         return instance;
     }
+    return factory.attach(await contractManager.getContract(name)) as MessageProxyForMainnetTester;
 }
diff --git a/test/utils/skale-manager-utils/contractManager.ts b/test/utils/skale-manager-utils/contractManager.ts
index 6c21724a2..5a71dd03e 100644
--- a/test/utils/skale-manager-utils/contractManager.ts
+++ b/test/utils/skale-manager-utils/contractManager.ts
@@ -13,31 +13,43 @@ export async function deployContractManager(contractManagerAddress: string) {
     if (contractManagerAddress === "0x0000000000000000000000000000000000000000") {
         instance = await contractManagerFactory.deploy() as ContractManager;
     } else {
-        instance = await contractManagerFactory.attach(contractManagerAddress) as ContractManager;
+        instance = contractManagerFactory.attach(contractManagerAddress) as ContractManager;
     }
-    if (await instance.getContract("KeyStorage") === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract("KeyStorage");
+    } catch {
         const keyStorageInstance = await (await ethers.getContractFactory("KeyStorageMock")).deploy() as KeyStorageMock;
         await instance.setContractsAddress("KeyStorage", keyStorageInstance);
     }
-    if (await instance.getContract(nameNodes) === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract(nameNodes);
+    } catch {
         const nodesInstance = await (await ethers.getContractFactory(nameNodes)).deploy() as Nodes;
         await instance.setContractsAddress(nameNodes, nodesInstance);
     }
-    if (await instance.getContract(nameSchains) === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract(nameSchains);
+    } catch {
         const schainsInstance = await (await ethers.getContractFactory(nameSchains)).deploy() as Schains;
         await schainsInstance.addContractManager(instance);
         await instance.setContractsAddress(nameSchains, schainsInstance);
     }
-    if (await instance.getContract(nameSchainsInternal) === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract(nameSchainsInternal);
+    } catch {
         const schainsInternalInstance = await (await ethers.getContractFactory(nameSchainsInternal)).deploy() as SchainsInternal;
         await schainsInternalInstance.addContractManager(instance);
         await instance.setContractsAddress(nameSchainsInternal, schainsInternalInstance);
     }
-    if (await instance.getContract(nameSkaleVerifier) === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract(nameSkaleVerifier);
+    } catch {
         const skaleVerifierInstance = await (await ethers.getContractFactory(nameSkaleVerifier)).deploy() as SkaleVerifierMock;
         await instance.setContractsAddress("SkaleVerifier", skaleVerifierInstance);
     }
-    if (await instance.getContract(nameWallets) === "0x0000000000000000000000000000000000000000") {
+    try {
+        await instance.getContract(nameWallets);
+     } catch {
         const walletsInstance = await (await ethers.getContractFactory(nameWallets)).deploy() as Wallets;
         await walletsInstance.addContractManager(instance);
         await instance.setContractsAddress(nameWallets, walletsInstance);