Add method to remove delegate with provided signature

gnosis/safe/api/transaction_service_api/transaction_service_api.py

@@ -3,7 +3,6 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlencode
 
-from eth_account.signers.local import LocalAccount
 from eth_typing import ChecksumAddress, Hash32, HexStr
 from hexbytes import HexBytes
 
@@ -265,39 +264,48 @@ def post_signatures(self, safe_tx_hash: bytes, signatures: bytes) -> bool:
 
     def add_delegate(
         self,
-        safe_address: ChecksumAddress,
         delegate_address: ChecksumAddress,
         delegator_address: ChecksumAddress,
         label: str,
-        signer_account: LocalAccount,
+        signature: bytes,
+        safe_address: Optional[ChecksumAddress] = None,
     ) -> bool:
-        hash_to_sign = self.create_delegate_message_hash(delegate_address)
-        signature = signer_account.signHash(hash_to_sign)
         add_payload = {
-            "safe": safe_address,
             "delegate": delegate_address,
             "delegator": delegator_address,
-            "signature": signature.signature.hex(),
+            "signature": HexBytes(signature).hex(),
             "label": label,
         }
+        if safe_address:
+            add_payload["safe"] = safe_address
         response = self._post_request("/api/v2/delegates/", add_payload)
         if not response.ok:
             raise SafeAPIException(f"Cannot add delegate: {response.content}")
         return True
 
     def remove_delegate(
         self,
-        safe_address: ChecksumAddress,
         delegate_address: ChecksumAddress,
-        signer_account: LocalAccount,
+        delegator_address: ChecksumAddress,
+        signature: bytes,
+        safe_address: Optional[ChecksumAddress] = None,
     ) -> bool:
-        hash_to_sign = self.create_delegate_message_hash(delegate_address)
-        signature = signer_account.signHash(hash_to_sign)
+        """
+        Deletes a delegated user
+
+        :param delegator_address:
+        :param delegate_address:
+        :param signature: Signature of a hash of an eip712 message.
+        :param safe_address: If specified, a delegate is removed for a delegator for the specific safe.
+            Otherwise, the delegate is deleted in a global form.
+        :return:
+        """
         remove_payload = {
-            "safe": safe_address,
-            "delegator": signer_account.address,
-            "signature": signature.signature.hex(),
+            "delegator": delegator_address,
+            "signature": HexBytes(signature).hex(),
         }
+        if safe_address:
+            remove_payload["safe"] = safe_address
         response = self._delete_request(
             f"/api/v2/delegates/{delegate_address}/",
             remove_payload,

gnosis/safe/tests/api/test_transaction_service_api.py

@@ -4,6 +4,7 @@
 
 from django.test import TestCase
 
+from eth_account import Account
 from eth_typing import HexStr
 from hexbytes import HexBytes
 
@@ -189,6 +190,80 @@ def test_get_safe_transaction(self):
                 str(context.exception),
             )
 
+    def test_remove_delegate(self):
+        with patch.object(
+            TransactionServiceApi, "_delete_request"
+        ) as mock_delete_request:
+            delegate_address = Account().create().address
+            delegator_account = Account().create()
+            message_hash = self.transaction_service_api.create_delegate_message_hash(
+                delegate_address
+            )
+            signature = delegator_account.signHash(message_hash).signature.hex()
+            self.transaction_service_api.remove_delegate(
+                delegate_address, delegator_account.address, signature
+            )
+
+            expected_url = f"/api/v2/delegates/{delegate_address}/"
+            expected_payload = {
+                "delegator": delegator_account.address,
+                "signature": signature,
+            }
+            mock_delete_request.assert_called_once_with(expected_url, expected_payload)
+
+            self.transaction_service_api.remove_delegate(
+                delegate_address,
+                delegator_account.address,
+                signature,
+                self.safe_address,
+            )
+
+            expected_payload = {
+                "safe": self.safe_address,
+                "delegator": delegator_account.address,
+                "signature": signature,
+            }
+            mock_delete_request.assert_called_with(expected_url, expected_payload)
+
+    def test_add_delegate(self):
+        with patch.object(TransactionServiceApi, "_post_request") as mock_post_request:
+            delegate_address = Account().create().address
+            delegator_account = Account().create()
+            message_hash = self.transaction_service_api.create_delegate_message_hash(
+                delegate_address
+            )
+            signature = delegator_account.signHash(message_hash).signature.hex()
+            label = "test label"
+            self.transaction_service_api.add_delegate(
+                delegate_address, delegator_account.address, label, signature
+            )
+
+            expected_url = "/api/v2/delegates/"
+            expected_payload = {
+                "delegate": delegate_address,
+                "delegator": delegator_account.address,
+                "signature": signature,
+                "label": label,
+            }
+            mock_post_request.assert_called_once_with(expected_url, expected_payload)
+
+            self.transaction_service_api.add_delegate(
+                delegate_address,
+                delegator_account.address,
+                label,
+                signature,
+                self.safe_address,
+            )
+
+            expected_payload = {
+                "safe": self.safe_address,
+                "delegate": delegate_address,
+                "delegator": delegator_account.address,
+                "signature": signature,
+                "label": label,
+            }
+            mock_post_request.assert_called_with(expected_url, expected_payload)
+
     def test_decode_data(self):
         with patch.object(TransactionServiceApi, "_post_request") as mock_post_request:
             mock_post_request.return_value.ok = True