Refactor almost duplicate functions

There were two very similar functions sign_balance_proof_close_message() and sign_balance_proof_update_message(). These two are combined into sign_balance_proof_message().
raiden-network · Aug 15, 2019 · e0d53ff · e0d53ff
1 parent 6816aa2
commit e0d53ff
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 36 deletions.
diff --git a/raiden_contracts/tests/fixtures/channel.py b/raiden_contracts/tests/fixtures/channel.py
@@ -6,7 +6,7 @@
 from web3 import Web3
 from web3.contract import Contract
 
-from raiden_contracts.constants import TEST_SETTLE_TIMEOUT_MIN, ChannelState
+from raiden_contracts.constants import TEST_SETTLE_TIMEOUT_MIN, ChannelState, MessageTypeId
 from raiden_contracts.tests.utils import (
     EMPTY_ADDITIONAL_HASH,
     EMPTY_BALANCE_HASH,
@@ -24,8 +24,7 @@
 from raiden_contracts.utils.proofs import (
     hash_balance_data,
     sign_balance_proof,
-    sign_balance_proof_close_message,
-    sign_balance_proof_update_message,
+    sign_balance_proof_message,
     sign_cooperative_settle_message,
     sign_withdraw_message,
 )
@@ -649,6 +648,7 @@ def get(
             _token_network.address,
             int(_token_network.functions.chain_id().call()),
             channel_identifier,
+            MessageTypeId.BALANCE_PROOF,
             balance_hash,
             nonce,
             additional_hash,
@@ -676,11 +676,12 @@ def get(
         _token_network = other_token_network or token_network
         private_key = get_private_key(participant)
 
-        non_closing_signature = sign_balance_proof_close_message(
+        non_closing_signature = sign_balance_proof_message(
             private_key,
             _token_network.address,
             int(_token_network.functions.chain_id().call()),
             channel_identifier,
+            MessageTypeId.BALANCE_PROOF,
             balance_hash,
             nonce,
             additional_hash,
@@ -709,11 +710,12 @@ def get(
         _token_network = other_token_network or token_network
         private_key = get_private_key(participant)
 
-        non_closing_signature = sign_balance_proof_update_message(
+        non_closing_signature = sign_balance_proof_message(
             private_key,
             _token_network.address,
             int(_token_network.functions.chain_id().call()),
             channel_identifier,
+            MessageTypeId.BALANCE_PROOF_UPDATE,
             balance_hash,
             nonce,
             additional_hash,
@@ -738,11 +740,12 @@ def get(
         _token_network = other_token_network or token_network
         private_key = get_private_key(participant)
 
-        non_closing_signature = sign_balance_proof_close_message(
+        non_closing_signature = sign_balance_proof_message(
             private_key,
             _token_network.address,
             int(_token_network.functions.chain_id().call()),
             channel_identifier,
+            MessageTypeId.BALANCE_PROOF,
             EMPTY_BALANCE_HASH,
             0,
             EMPTY_ADDITIONAL_HASH,

diff --git a/raiden_contracts/utils/proofs.py b/raiden_contracts/utils/proofs.py
@@ -128,6 +128,7 @@ def sign_balance_proof(
     token_network_address: HexAddress,
     chain_identifier: int,
     channel_identifier: int,
+    msg_type: MessageTypeId,
     balance_hash: bytes,
     nonce: int,
     additional_hash: bytes,
@@ -141,45 +142,19 @@ def sign_balance_proof(
             balance_hash=balance_hash,
             nonce=nonce,
             additional_hash=additional_hash,
-            msg_type=MessageTypeId.BALANCE_PROOF,
-        )
-    )
-
-    return sign(privkey=privatekey, msg_hash=message_hash, v=v)
-
-
-def sign_balance_proof_close_message(
-    privatekey: str,
-    token_network_address: HexAddress,
-    chain_identifier: int,
-    channel_identifier: int,
-    balance_hash: bytes,
-    nonce: int,
-    additional_hash: bytes,
-    closing_signature: bytes,
-    v: int = 27,
-) -> bytes:
-    message_hash = eth_sign_hash_message(
-        pack_balance_proof_message(
-            token_network_address=token_network_address,
-            chain_identifier=chain_identifier,
-            channel_identifier=channel_identifier,
-            msg_type=MessageTypeId.BALANCE_PROOF,
-            balance_hash=balance_hash,
-            nonce=nonce,
-            additional_hash=additional_hash,
-            closing_signature=closing_signature,
+            msg_type=msg_type,
         )
     )
 
     return sign(privkey=privatekey, msg_hash=message_hash, v=v)
 
 
-def sign_balance_proof_update_message(
+def sign_balance_proof_message(
     privatekey: str,
     token_network_address: HexAddress,
     chain_identifier: int,
     channel_identifier: int,
+    msg_type: MessageTypeId,
     balance_hash: bytes,
     nonce: int,
     additional_hash: bytes,
@@ -191,7 +166,7 @@ def sign_balance_proof_update_message(
             token_network_address=token_network_address,
             chain_identifier=chain_identifier,
             channel_identifier=channel_identifier,
-            msg_type=MessageTypeId.BALANCE_PROOF_UPDATE,
+            msg_type=msg_type,
             balance_hash=balance_hash,
             nonce=nonce,
             additional_hash=additional_hash,