Fix bugs in SafeOperator and SafeTxServiceOperator

.pre-commit-config.yaml

@@ -2,11 +2,11 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/PyCQA/isort
-    rev: 6.0.1
+    rev: 8.0.1
     hooks:
       - id: isort
   - repo: https://github.com/psf/black
-    rev: 25.9.0
+    rev: 26.1.0
     hooks:
       - id: black
   - repo: https://github.com/PyCQA/flake8

docker-compose.yml

@@ -0,0 +1,15 @@
+services:
+  safe-cli:
+    build:
+      context: .
+      dockerfile: ./Dockerfile
+
+  ganache:
+    image: trufflesuite/ganache:latest
+    ports:
+      - "8545:8545"
+    command: --defaultBalanceEther 10000 --gasLimit 10000000 -a 10 --chain.chainId 1337 --chain.networkId 1337 -d --host 0.0.0.0
+    healthcheck:
+      test: bash -c "echo 'hello' > /dev/tcp/localhost/8545"
+    profiles:
+      - develop

run_tests.sh

@@ -1,3 +1,5 @@
 #/usr/bin/env sh
+docker compose --profile develop up -d ganache
 pip install -e .
 pytest
+docker compose down --profile develop

src/safe_cli/operators/safe_operator.py

@@ -744,6 +744,8 @@ def update_version_to_l2(
                     safe_l2_singleton[0], fallback_handler[0]
                 ).build_transaction(get_empty_tx_params())["data"]
             )
+            safe_l2_singleton = safe_l2_singleton[0]
+            fallback_handler = fallback_handler[0]
         elif safe_version in ("1.3.0", "1.4.1"):
             safe_l2_singleton = safe_deployments[safe_version]["GnosisSafeL2"][
                 str(chain_id)
@@ -754,6 +756,7 @@ def update_version_to_l2(
                     safe_l2_singleton[0]
                 ).build_transaction(get_empty_tx_params())["data"]
             )
+            safe_l2_singleton = safe_l2_singleton[0]
         else:
             raise InvalidMasterCopyException(
                 "Current version is not supported to migrate to L2"
@@ -1215,7 +1218,7 @@ def drain(self, to: str):
             safe_txs.append(safe_tx)
 
         if safe_txs:
-            multisend_tx = self.batch_safe_txs(self.get_nonce(), safe_txs)
+            multisend_tx = self.batch_safe_txs(self.safe.retrieve_nonce(), safe_txs)
             if multisend_tx is not None:
                 if self.execute_safe_transaction(multisend_tx):
                     print_formatted_text(

src/safe_cli/operators/safe_tx_service_operator.py

@@ -102,6 +102,7 @@ def confirm_message(self, safe_message_hash: bytes, sender: ChecksumAddress):
                     f"<ansired>Message with hash {to_0x_hex_str(safe_message_hash)} does not exist</ansired>"
                 )
             )
+            return False
         if not yes_or_no_question(
             f"Message: {safe_message['message']} \n Do you want to sign the following message?:"
         ):
@@ -300,13 +301,15 @@ def execute_tx(self, safe_tx_hash: Sequence[bytes]) -> bool:
                     f"has already been executed on {to_0x_hex_str(tx_hash)}</ansired>"
                 )
             )
+            return False
         elif len(safe_tx.signers) < self.safe_cli_info.threshold:
             print_formatted_text(
                 HTML(
                     f"<ansired>Number of signatures {len(safe_tx.signers)} "
                     f"must reach the threshold {self.safe_cli_info.threshold}</ansired>"
                 )
             )
+            return False
         else:
             if executed := self.execute_safe_transaction(safe_tx):
                 self.refresh_safe_cli_info()
@@ -416,18 +419,18 @@ def get_permitted_signers(self) -> Set[ChecksumAddress]:
     def drain(self, to: ChecksumAddress):
         balances = self.safe_tx_service.get_balances(self.address)
         safe_txs = []
-        safe_tx = None
         for balance in balances:
             amount = int(balance["balance"])
+            if amount == 0:
+                continue
             if balance["tokenAddress"] is None:  # Then is ether
-                if amount != 0:
-                    safe_tx = self.prepare_safe_transaction(
-                        to,
-                        amount,
-                        b"",
-                        SafeOperationEnum.CALL,
-                        safe_nonce=None,
-                    )
+                safe_tx = self.prepare_safe_transaction(
+                    to,
+                    amount,
+                    b"",
+                    SafeOperationEnum.CALL,
+                    safe_nonce=None,
+                )
             else:
                 transaction = (
                     get_erc20_contract(self.ethereum_client.w3, balance["tokenAddress"])
@@ -441,10 +444,9 @@ def drain(self, to: ChecksumAddress):
                     SafeOperationEnum.CALL,
                     safe_nonce=None,
                 )
-            if safe_tx:
-                safe_txs.append(safe_tx)
+            safe_txs.append(safe_tx)
         if len(safe_txs) > 0:
-            multisend_tx = self.batch_safe_txs(safe_tx.safe_nonce, safe_txs)
+            multisend_tx = self.batch_safe_txs(safe_txs[0].safe_nonce, safe_txs)
             if multisend_tx is not None:
                 self.post_transaction_to_tx_service(multisend_tx)
                 print_formatted_text(

tests/test_safe_operator.py

@@ -393,6 +393,19 @@ def test_update_to_l2_v111(self):
                 safe_operator_v111.safe.retrieve_fallback_handler(),
                 self.compatibility_fallback_handler_V1_3_0.address,
             )
+            # safe_cli_info must cache addresses (strings), not the raw deployment lists
+            self.assertIsInstance(safe_operator_v111.safe_cli_info.master_copy, str)
+            self.assertEqual(
+                safe_operator_v111.safe_cli_info.master_copy,
+                self.safe_contract_V1_3_0.address,
+            )
+            self.assertIsInstance(
+                safe_operator_v111.safe_cli_info.fallback_handler, str
+            )
+            self.assertEqual(
+                safe_operator_v111.safe_cli_info.fallback_handler,
+                self.compatibility_fallback_handler_V1_3_0.address,
+            )
 
     def test_update_to_l2_v130(self):
         migration_contract_address = self._deploy_l2_migration_contract()
@@ -430,6 +443,12 @@ def test_update_to_l2_v130(self):
                 safe_operator_v130.safe.retrieve_fallback_handler(),
                 previous_fallback_handler,
             )
+            # safe_cli_info must cache an address (string), not the raw deployment list
+            self.assertIsInstance(safe_operator_v130.safe_cli_info.master_copy, str)
+            self.assertEqual(
+                safe_operator_v130.safe_cli_info.master_copy,
+                safe_contract_l2_130_address,
+            )
 
     def test_drain(self):
         safe_operator = self.setup_operator()

tests/test_safe_tx_service_operator.py

@@ -602,9 +602,83 @@ def test_confirm_message(
             safe_operator.confirm_message(safe_message_hash, sender.address)
         )
 
-    def test_drain(self):
-        # TODO Drain is a complex to mock
-        pass
+        # When get_message itself raises SafeAPIException, confirm_message must return
+        # False instead of raising NameError (safe_message would be unbound)
+        post_message_signature.side_effect = None
+        get_message.side_effect = SafeAPIException("Message not found")
+        self.assertFalse(
+            safe_operator.confirm_message(safe_message_hash, sender.address)
+        )
+
+    @mock.patch.object(
+        SafeTx, "safe_version", return_value="1.4.1", new_callable=mock.PropertyMock
+    )
+    @mock.patch.object(TransactionServiceApi, "_get_request")
+    def test_execute_tx_returns_false(
+        self,
+        _get_request_mock: MagicMock,
+        safe_version_mock: mock.PropertyMock,
+    ):
+        safe_operator = self.setup_operator(
+            number_owners=1, mode=SafeOperatorMode.TX_SERVICE
+        )
+        safe_tx_hash = HexBytes(
+            "0xae1c18dd9fca652b83743fc0b0ac2d396c68d523b49f4d41af5b00dc2f995bf6"
+        )
+
+        # Already-executed tx must return False, not None
+        _get_request_mock.return_value = GetMultisigTxRequestMock(executed=True)
+        self.assertIs(safe_operator.execute_tx(safe_tx_hash), False)
+
+        # Not enough signatures must return False, not None.
+        # Force threshold above the single signer present in the mock.
+        safe_operator.safe_cli_info.threshold = 2
+        _get_request_mock.return_value = GetMultisigTxRequestMock(executed=False)
+        self.assertIs(safe_operator.execute_tx(safe_tx_hash), False)
+
+    @mock.patch.object(TransactionServiceApi, "get_balances")
+    @mock.patch("safe_cli.operators.safe_tx_service_operator.get_erc20_contract")
+    def test_drain(
+        self, get_erc20_contract_mock: MagicMock, get_balances_mock: MagicMock
+    ):
+        safe_operator = self.setup_operator(
+            number_owners=1, mode=SafeOperatorMode.TX_SERVICE
+        )
+        to = Account.create().address
+        token_address = "0x6B175474E89094C44Da98b954EedeAC495271d0F"
+
+        # Token balance followed by zero ether balance.  The previous bug caused the
+        # token safe_tx to be appended twice because the loop variable was not reset
+        # when ether amount == 0.
+        get_balances_mock.return_value = [
+            {"tokenAddress": token_address, "balance": "1000000000000000000"},
+            {"tokenAddress": None, "balance": "0"},
+        ]
+        mock_contract = MagicMock()
+        mock_contract.functions.transfer.return_value.build_transaction.return_value = {
+            "data": b"\x00"
+        }
+        get_erc20_contract_mock.return_value = mock_contract
+
+        safe_tx_mock = MagicMock()
+        safe_tx_mock.safe_nonce = 0
+        batch_safe_txs_mock = MagicMock(return_value=None)
+
+        with (
+            mock.patch.object(
+                safe_operator, "prepare_safe_transaction", return_value=safe_tx_mock
+            ),
+            mock.patch.object(safe_operator, "batch_safe_txs", batch_safe_txs_mock),
+        ):
+            safe_operator.drain(to)
+
+        batch_safe_txs_mock.assert_called_once()
+        _, safe_txs_passed = batch_safe_txs_mock.call_args[0]
+        self.assertEqual(
+            len(safe_txs_passed),
+            1,
+            "drain must not duplicate transactions when ether balance is zero",
+        )
 
 
 if __name__ == "__main__":