Refactor MultiSend usage

ademidun · Jul 18, 2022 · 5b53f56 · 5b53f56
1 parent 96e7311
commit 5b53f56
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 43 deletions.
diff --git a/safe_cli/operators/safe_operator.py b/safe_cli/operators/safe_operator.py
@@ -28,12 +28,7 @@
 from safe_cli.api.relay_service_api import RelayServiceApi
 from safe_cli.api.transaction_service_api import TransactionServiceApi
 from safe_cli.ethereum_hd_wallet import get_account_from_words
-from safe_cli.safe_addresses import (
-    LAST_DEFAULT_CALLBACK_HANDLER,
-    LAST_MULTISEND_CALL_ONLY_CONTRACT,
-    LAST_MULTISEND_CONTRACT,
-    LAST_SAFE_CONTRACT,
-)
+from safe_cli.safe_addresses import LAST_DEFAULT_CALLBACK_HANDLER, LAST_SAFE_CONTRACT
 from safe_cli.utils import get_erc_20_list, yes_or_no_question
 
 try:
@@ -530,7 +525,7 @@ def update_version(self) -> Optional[bool]:
                 "Not valid addresses to update Safe", *addresses
             )
 
-        multisend = MultiSend(LAST_MULTISEND_CONTRACT, self.ethereum_client)
+        multisend = MultiSend(ethereum_client=self.ethereum_client)
         tx_params = {"from": self.address, "gas": 0, "gasPrice": 0}
         multisend_txs = [
             MultiSendTx(MultiSendOperation.CALL, self.address, 0, data)
@@ -763,11 +758,14 @@ def batch_safe_txs(self, safe_nonce: int, safe_txs: [SafeTx]) -> SafeTx:
 
         :return:
         """
-        if not self.ethereum_client.is_contract(LAST_MULTISEND_CALL_ONLY_CONTRACT):
+
+        try:
+            multisend = MultiSend(ethereum_client=self.ethereum_client)
+        except ValueError:
             print_formatted_text(
                 HTML(
-                    f"<ansired>Multisend call only contract {LAST_MULTISEND_CALL_ONLY_CONTRACT} "
-                    f"is not deployed on this network and it's required for batching txs</ansired>"
+                    "<ansired>Multisend contract is not deployed on this network and it's required for "
+                    "batching txs</ansired>"
                 )
             )
 
@@ -785,13 +783,10 @@ def batch_safe_txs(self, safe_nonce: int, safe_txs: [SafeTx]) -> SafeTx:
                 )
 
         if len(multisend_txs) > 1:
-            multisend = MultiSend(
-                LAST_MULTISEND_CALL_ONLY_CONTRACT, self.ethereum_client
-            )
             safe_tx = SafeTx(
                 self.ethereum_client,
                 self.address,
-                LAST_MULTISEND_CALL_ONLY_CONTRACT,
+                multisend.address,
                 0,
                 multisend.build_tx_data(multisend_txs),
                 SafeOperation.DELEGATE_CALL.value,

diff --git a/safe_cli/operators/safe_tx_service_operator.py b/safe_cli/operators/safe_tx_service_operator.py
@@ -12,7 +12,6 @@
 from safe_cli.api.base_api import BaseAPIException
 from safe_cli.utils import yes_or_no_question
 
-from ..safe_addresses import LAST_MULTISEND_CALL_ONLY_CONTRACT
 from .safe_operator import (
     AccountNotLoadedException,
     NonExistingOwnerException,
@@ -125,11 +124,13 @@ def batch_txs(self, safe_nonce: int, safe_tx_hashes: Sequence[bytes]) -> bool:
         :return:
         """
 
-        if not self.ethereum_client.is_contract(LAST_MULTISEND_CALL_ONLY_CONTRACT):
+        try:
+            multisend = MultiSend(ethereum_client=self.ethereum_client)
+        except ValueError:
             print_formatted_text(
                 HTML(
-                    f"<ansired>Multisend call only contract {LAST_MULTISEND_CALL_ONLY_CONTRACT} "
-                    f"is not deployed on this network and it's required for batching txs</ansired>"
+                    "<ansired>Multisend contract is not deployed on this network and it's required for "
+                    "batching txs</ansired>"
                 )
             )
 
@@ -148,13 +149,10 @@ def batch_txs(self, safe_nonce: int, safe_tx_hashes: Sequence[bytes]) -> bool:
                 )
 
         if len(multisend_txs) > 1:
-            multisend = MultiSend(
-                LAST_MULTISEND_CALL_ONLY_CONTRACT, self.ethereum_client
-            )
             safe_tx = SafeTx(
                 self.ethereum_client,
                 self.address,
-                LAST_MULTISEND_CALL_ONLY_CONTRACT,
+                multisend.address,
                 0,
                 multisend.build_tx_data(multisend_txs),
                 SafeOperation.DELEGATE_CALL.value,

diff --git a/tests/test_safe_operator.py b/tests/test_safe_operator.py
@@ -246,31 +246,31 @@ def test_drain(self):
         _, _, contract_address = MultiSend.deploy_contract(
             self.ethereum_client, account
         )
-        # Getting events filtered by Transfer
-        last = safe_operator.ethereum_client.get_block("latest")["number"]
-        token_address = get_erc_20_list(
-            safe_operator.ethereum_client, safe_operator.address, 1, last
-        )
-        self.assertEqual(len(token_address), num_contracts_erc20)
-        for token in token_address:
-            result = self.ethereum_client.erc20.get_balance(
-                safe_operator.address, token
+        with mock.patch.object(MultiSend, "MULTISEND_ADDRESSES", [contract_address]):
+            # Getting events filtered by Transfer
+            last = safe_operator.ethereum_client.get_block("latest")["number"]
+            token_address = get_erc_20_list(
+                safe_operator.ethereum_client, safe_operator.address, 1, last
             )
-            self.assertEqual(result, num_transactions)
-        # Draining the account to a new account
-        with mock.patch(
-            "safe_cli.operators.safe_operator.LAST_MULTISEND_CALL_ONLY_CONTRACT",
-            contract_address,
-        ):
+            self.assertEqual(len(token_address), num_contracts_erc20)
+            for token in token_address:
+                result = self.ethereum_client.erc20.get_balance(
+                    safe_operator.address, token
+                )
+                self.assertEqual(result, num_transactions)
+            # Draining the account to a new account
             safe_operator.drain(account.address)
-        # Checking that the account is empty
-        for token in token_address:
+            # Checking that the account is empty
+            for token in token_address:
+                self.assertEqual(
+                    self.ethereum_client.erc20.get_balance(
+                        safe_operator.address, token
+                    ),
+                    0,
+                )
             self.assertEqual(
-                self.ethereum_client.erc20.get_balance(safe_operator.address, token), 0
+                safe_operator.ethereum_client.get_balance(safe_operator.address), 0
             )
-        self.assertEqual(
-            safe_operator.ethereum_client.get_balance(safe_operator.address), 0
-        )
 
 
 if __name__ == "__main__":