Get private Safes balances from blockchain

- Closes #125
5afe · Jul 19, 2019 · 378d303 · 378d303
1 parent b87bd00
commit 378d303
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 19 deletions.
diff --git a/safe_relay_service/relay/serializers.py b/safe_relay_service/relay/serializers.py
@@ -17,6 +17,8 @@
                                              EthereumTxType, InternalTx,
                                              SafeFunding)
 
+from .services import StatsServiceProvider
+
 logger = logging.getLogger(__name__)
 
 
@@ -84,8 +86,10 @@ class TokensWithBalanceSerializer(serializers.Serializer):
 class SafeContractSerializer(serializers.Serializer):
     created = serializers.DateTimeField()
     address = EthereumAddressField()
-    balance = serializers.IntegerField(min_value=0, allow_null=True)
-    tokens_with_balance = TokensWithBalanceSerializer(source='get_tokens_with_balance', many=True)
+    tokens_with_balance = serializers.SerializerMethodField()
+
+    def get_tokens_with_balance(self, obj):
+        return StatsServiceProvider().get_balances(obj.address)
 
 
 class SafeBalanceResponseSerializer(serializers.Serializer):
@@ -165,15 +169,15 @@ def get_fields(self):
         result['from'] = _from
         return result
 
+    def get_tx_type(self, obj) -> str:
+        return EthereumTxType(obj.tx_type).name
+
     def get_call_type(self, obj) -> str:
         if obj.call_type is None:
             return None
         else:
             return EthereumTxCallType(obj.call_type).name
 
-    def get_tx_type(self, obj) -> str:
-        return EthereumTxType(obj.tx_type).name
-
 
 class EthereumTxSerializer(serializers.ModelSerializer):
     class Meta:

diff --git a/safe_relay_service/relay/tests/test_views.py b/safe_relay_service/relay/tests/test_views.py
@@ -825,34 +825,36 @@ def test_private_safes_view(self):
         self.assertEqual(dateparse.parse_datetime(safe_response['created']), safe_contract.created)
 
         # Add balance
-        InternalTxFactory(to=safe_contract.address, value=10)
-        InternalTxFactory(to=safe_contract.address, value=5)
-        InternalTxFactory(_from=safe_contract.address, value=8)
+        ether_balance = 7
+        self.send_ether(safe_contract.address, ether_balance)
         # Balance of safe_contract must be 15 - 8 = 7 now
         response = self.client.get(url)
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.json()['count'], 1)
         results = response.json()['results']
         safe_response = results[0]
         self.assertEqual(safe_response['address'], safe_contract.address)
-        self.assertEqual(safe_response['balance'], 7)
+        print(safe_response)
+        self.assertIsNone(safe_response['tokensWithBalance'][0]['tokenAddress'])
+        self.assertEqual(safe_response['tokensWithBalance'][0]['balance'], ether_balance)
 
         # Add token transfers
-        ethereum_event = EthereumEventFactory(to=safe_contract.address, value=19)
-        EthereumEventFactory(from_=safe_contract.address, token_address=ethereum_event.token_address, value=5)
-        ethereum_event_2 = EthereumEventFactory(to=safe_contract.address, value=4)
-        EthereumEventFactory(from_=safe_contract.address, token_address=ethereum_event_2.token_address, value=5)
+        token_balance = 8
+        example_erc20_1 = self.deploy_example_erc20(token_balance, safe_contract.address)
+        example_erc20_2 = self.deploy_example_erc20(token_balance, safe_contract.address)
+        EthereumEventFactory(to=safe_contract.address, value=19, token_address=example_erc20_1.address)
+        EthereumEventFactory(to=safe_contract.address, value=4, token_address=example_erc20_2.address)
         response = self.client.get(url)
         self.assertEqual(response.status_code, status.HTTP_200_OK)
         self.assertEqual(response.json()['count'], 1)
         results = response.json()['results']
         safe_response = results[0]
         self.assertEqual(safe_response['address'], safe_contract.address)
-        self.assertEqual(safe_response['balance'], 7)
-        self.assertEqual(safe_response['tokensWithBalance'][0]['tokenAddress'], ethereum_event.token_address)
-        self.assertEqual(safe_response['tokensWithBalance'][0]['balance'], 14)
-        self.assertEqual(safe_response['tokensWithBalance'][1]['tokenAddress'], ethereum_event_2.token_address)
-        self.assertEqual(safe_response['tokensWithBalance'][1]['balance'], -1)
+        self.assertCountEqual(safe_response['tokensWithBalance'], [
+            {'tokenAddress': None, 'balance': ether_balance},
+            {'tokenAddress': example_erc20_1.address, 'balance': token_balance},
+            {'tokenAddress': example_erc20_2.address, 'balance': token_balance},
+        ])
         self.client.force_authenticate(user=None)
 
     def test_api_token_auth(self):

diff --git a/safe_relay_service/relay/views.py b/safe_relay_service/relay/views.py
@@ -505,5 +505,5 @@ def get(self, request, format=None):
 class PrivateSafesView(ListAPIView):
     authentication_classes = (TokenAuthentication,)
     permission_classes = (IsAuthenticated,)
-    queryset = SafeContract.objects.deployed().with_balance().order_by('created')
+    queryset = SafeContract.objects.deployed().order_by('created')
     serializer_class = SafeContractSerializer