Repository: rotki/rotki

Issue:
Assertion in the logs for test user
## Problem Definition

The following assertion was seen at the logs for a test user

```
04/08/2019 22:14:53 CEST -- ERROR:rotkehlchen.server:Greenlet for task 1 dies with exception: at this point the key should only be Asset type.                                                                                                                                   
Exception Name: <class 'AssertionError'>                                                                                                
Exception Info: at this point the key should only be Asset type                                                                         
Traceback:                                                                                                                              
   File "src/gevent/greenlet.py", line 716, in gevent._greenlet.Greenlet.run
  File "/home/lefteris/w/rotkehlchen/rotkehlchen/server.py", line 128, in _query_async
    result = getattr(self, command)(**kwargs)
  File "/home/lefteris/w/rotkehlchen/rotkehlchen/server.py", line 326, in query_balances                                 
    result = self.rotkehlchen.query_balances(save_data)                                                                                 
  File "/home/lefteris/w/rotkehlchen/rotkehlchen/rotkehlchen.py", line 483, in query_balances                          
    self.data.save_balances_data(data=result_dict, timestamp=timestamp)                                                           
  File "/home/lefteris/w/rotkehlchen/rotkehlchen/data_handler.py", line 228, in save_balances_data                               
    self.db.write_balances_data(data=data, timestamp=timestamp)                                                                         
  File "/home/lefteris/w/rotkehlchen/rotkehlchen/db/dbhandler.py", line 658, in write_balances_data                               
    assert isinstance(key, Asset), 'at this point the key should only be Asset type'         
```

## Task

- Find out why it happened and fix it
- Write a test for it



Discussion:

Resolving patch:
diff --git a/rotkehlchen/data_handler.py b/rotkehlchen/data_handler.py
--- a/rotkehlchen/data_handler.py
+++ b/rotkehlchen/data_handler.py
@@ -367,7 +367,7 @@ def set_fiat_balance(
 
         return True, ''
 
-    def get_fiat_balances(self) -> Dict[FiatAsset, str]:
+    def get_fiat_balances(self) -> Dict[Asset, str]:
         return self.db.get_fiat_balances()
 
     def get_external_trades(
diff --git a/rotkehlchen/db/dbhandler.py b/rotkehlchen/db/dbhandler.py
--- a/rotkehlchen/db/dbhandler.py
+++ b/rotkehlchen/db/dbhandler.py
@@ -600,7 +600,7 @@ def remove_fiat_balance(self, currency: Asset) -> None:
         self.conn.commit()
         self.update_last_write()
 
-    def get_fiat_balances(self) -> Dict[FiatAsset, str]:
+    def get_fiat_balances(self) -> Dict[Asset, str]:
         cursor = self.conn.cursor()
         query = cursor.execute(
             'SELECT asset, amount FROM current_balances;',
@@ -608,7 +608,13 @@ def get_fiat_balances(self) -> Dict[FiatAsset, str]:
 
         result = {}
         for entry in query:
-            result[entry[0]] = entry[1]
+            try:
+                result[Asset(entry[0])] = entry[1]
+            except UnknownAsset:
+                self.msg_aggregator.add_error(
+                    f'Unknown FIAT asset {entry[0]} found in the DB. Skipping it ...',
+                )
+
         return result
 
     def get_blockchain_accounts(self) -> BlockchainAccounts:
diff --git a/rotkehlchen/rotkehlchen.py b/rotkehlchen/rotkehlchen.py
--- a/rotkehlchen/rotkehlchen.py
+++ b/rotkehlchen/rotkehlchen.py
@@ -392,13 +392,12 @@ def process_history(self, start_ts, end_ts):
         )
         return result, error_or_empty
 
-    def query_fiat_balances(self):
-        log.info('query_fiat_balances called')
+    def query_fiat_balances(self) -> Dict[Asset, Dict[str, FVal]]:
         result = {}
         balances = self.data.get_fiat_balances()
         for currency, amount in balances.items():
             amount = FVal(amount)
-            usd_rate = Inquirer().query_fiat_pair(currency, 'USD')
+            usd_rate = Inquirer().query_fiat_pair(currency, S_USD)
             result[currency] = {
                 'amount': amount,
                 'usd_value': amount * usd_rate,


Test patch:
diff --git a/rotkehlchen/tests/db/test_db.py b/rotkehlchen/tests/db/test_db.py
--- a/rotkehlchen/tests/db/test_db.py
+++ b/rotkehlchen/tests/db/test_db.py
@@ -11,7 +11,7 @@
 
 from rotkehlchen.assets.asset import Asset, EthereumToken
 from rotkehlchen.constants import YEAR_IN_SECONDS
-from rotkehlchen.constants.assets import A_BTC, A_ETH, A_USD, FIAT_CURRENCIES, S_CNY, S_EUR
+from rotkehlchen.constants.assets import A_BTC, A_ETH, A_EUR, A_USD, FIAT_CURRENCIES, S_CNY, S_EUR
 from rotkehlchen.data_handler import DataHandler
 from rotkehlchen.db.dbhandler import (
     DBINFO_FILENAME,
@@ -34,7 +34,7 @@
 )
 from rotkehlchen.errors import AuthenticationError, InputError
 from rotkehlchen.fval import FVal
-from rotkehlchen.tests.utils.constants import A_DAO, A_DOGE, A_GNO, A_RDN, A_XMR
+from rotkehlchen.tests.utils.constants import A_CNY, A_DAO, A_DOGE, A_GNO, A_RDN, A_XMR
 from rotkehlchen.tests.utils.rotkehlchen import add_starting_balances
 from rotkehlchen.typing import SupportedBlockchain, Timestamp
 from rotkehlchen.user_messages import MessagesAggregator
@@ -100,7 +100,7 @@ def test_export_import_db(data_dir, username):
     data.decompress_and_decrypt_db('123', encoded_data)
     fiat_balances = data.get_fiat_balances()
     assert len(fiat_balances) == 1
-    assert int(fiat_balances['EUR']) == 10
+    assert int(fiat_balances[A_EUR]) == 10
 
 
 def test_writting_fetching_data(data_dir, username):
@@ -497,13 +497,13 @@ def test_data_set_fiat_balance(data_dir, username):
     assert success
     balances = data.get_fiat_balances()
     assert len(balances) == 2
-    assert balances[S_EUR] == amount_eur
-    assert balances[S_CNY] == amount_cny
+    assert balances[A_EUR] == amount_eur
+    assert balances[A_CNY] == amount_cny
 
     success, _ = data.set_fiat_balance(S_EUR, '')
     balances = data.get_fiat_balances()
     assert len(balances) == 1
-    assert balances[S_CNY] == amount_cny
+    assert balances[A_CNY] == amount_cny
 
     # also check that all the fiat assets in the fiat table are in
     # all_assets.json
diff --git a/rotkehlchen/tests/test_api.py b/rotkehlchen/tests/test_api.py
--- a/rotkehlchen/tests/test_api.py
+++ b/rotkehlchen/tests/test_api.py
@@ -4,13 +4,17 @@
 
 import pytest
 
+from rotkehlchen.assets.asset import Asset
+from rotkehlchen.constants.assets import A_BTC, A_ETH
 from rotkehlchen.data_handler import VALID_SETTINGS
 from rotkehlchen.db.utils import ROTKEHLCHEN_DB_VERSION
 from rotkehlchen.fval import FVal
 from rotkehlchen.inquirer import Inquirer
-from rotkehlchen.tests.utils.mock import MockWeb3
+from rotkehlchen.tests.utils.exchanges import BINANCE_BALANCES_RESPONSE
+from rotkehlchen.tests.utils.mock import MockResponse, MockWeb3
 from rotkehlchen.tests.utils.rotkehlchen import add_starting_balances
-from rotkehlchen.typing import Timestamp
+from rotkehlchen.typing import Price, Timestamp
+from rotkehlchen.utils.misc import ts_now
 from rotkehlchen.utils.serialization import rlk_jsonloads_dict
 
 
@@ -204,6 +208,93 @@ def test_query_timed_balances_data(rotkehlchen_server):
     assert response['message'] == 'Unknown asset DSXXADA provided.'
 
 
+def test_query_balances(rotkehlchen_server, function_scope_binance):
+    """Test that the query_balances call works properly.
+
+    That means that the balances are all returned in the expected format and
+    that they are saved in the DB.
+
+    The test is for a user with fiat balances and with some binance balances.
+    """
+
+    # First set the fiat balances
+    ok, _ = rotkehlchen_server.set_fiat_balance('USD', '100.5')
+    assert ok
+    ok, _ = rotkehlchen_server.set_fiat_balance('EUR', '75.5')
+    assert ok
+    rotkehlchen_server.rotkehlchen.binance = function_scope_binance
+    rotkehlchen_server.rotkehlchen.connected_exchanges.append('binance')
+
+    def mock_binance_balances(url):  # pylint: disable=unused-argument
+        return MockResponse(200, BINANCE_BALANCES_RESPONSE)
+
+    mock_binance = patch.object(
+        rotkehlchen_server.rotkehlchen.binance.session,
+        'get',
+        side_effect=mock_binance_balances,
+    )
+
+    eur_usd_rate = Inquirer().query_fiat_pair('EUR', 'USD')
+
+    eth_usd_rate = FVal('100.5')
+    btc_usd_rate = FVal('120.1')
+
+    def mock_query_cryptocompare_for_fiat_price(asset: Asset) -> Price:
+        if asset == A_ETH:
+            return Price(eth_usd_rate)
+        elif asset == A_BTC:
+            return Price(btc_usd_rate)
+
+        # else
+        raise AssertionError(f'Unexpected asset {asset} at mock cryptocompare query')
+
+    mock_find_usd_price = patch(
+        'rotkehlchen.inquirer.query_cryptocompare_for_fiat_price',
+        side_effect=mock_query_cryptocompare_for_fiat_price,
+    )
+
+    now = ts_now()
+    with mock_binance, mock_find_usd_price:
+        result = rotkehlchen_server.query_balances(save_data=True)
+
+    assert result['USD']['amount'] == '100.5'
+    assert result['USD']['usd_value'] == '100.5'
+    eur_amount = FVal('75.5')
+    assert result['EUR']['amount'] == str(eur_amount)
+    eur_usd_value = eur_amount * eur_usd_rate
+    assert result['EUR']['usd_value'] == str(eur_usd_value)
+    eth_amount = FVal('4763368.68006011')
+    assert result['ETH']['amount'] == str(eth_amount)
+    eth_usd_value = eth_amount * eth_usd_rate
+    assert result['ETH']['usd_value'] == str(eth_usd_value)
+    btc_amount = FVal('4723846.89208129')
+    assert result['BTC']['amount'] == str(btc_amount)
+    btc_usd_value = btc_amount * btc_usd_rate
+    assert result['BTC']['usd_value'] == str(btc_usd_value)
+    binance_usd_value = btc_usd_value + eth_usd_value
+    assert result['location']['binance']['usd_value'] == str(binance_usd_value)
+    banks_usd_value = eur_usd_value + FVal('100.5')
+    assert result['location']['banks']['usd_value'] == str(banks_usd_value)
+    assert result['net_usd'] == str(banks_usd_value + binance_usd_value)
+
+    # make sure that balances also got saved in the DB
+    db = rotkehlchen_server.rotkehlchen.data.db
+    save_ts = db.get_last_balance_save_time()
+    assert save_ts >= now
+    assert save_ts - now < 5, 'Saving balances took too long'
+
+    location_data = db.get_latest_location_value_distribution()
+    assert len(location_data) == 4
+
+    assert location_data[0].location == 'banks'
+    assert location_data[0].usd_value == str(banks_usd_value)
+    assert location_data[1].location == 'binance'
+    assert location_data[1].usd_value == str(binance_usd_value)
+    assert location_data[2].location == 'blockchain'
+    assert location_data[3].location == 'total'
+    assert location_data[3].usd_value == str(banks_usd_value + binance_usd_value)
+
+
 def test_query_netvalue_data(rotkehlchen_server):
     """Test that query_netvalue_data API call works as expected"""
     add_starting_balances(rotkehlchen_server.rotkehlchen.data)
diff --git a/rotkehlchen/tests/test_binance.py b/rotkehlchen/tests/test_binance.py
--- a/rotkehlchen/tests/test_binance.py
+++ b/rotkehlchen/tests/test_binance.py
@@ -17,6 +17,7 @@
 from rotkehlchen.fval import FVal
 from rotkehlchen.order_formatting import Trade, TradeType
 from rotkehlchen.tests.utils.constants import A_BNB, A_RDN, A_USDT
+from rotkehlchen.tests.utils.exchanges import BINANCE_BALANCES_RESPONSE
 from rotkehlchen.tests.utils.factories import make_random_b64bytes
 from rotkehlchen.tests.utils.mock import MockResponse
 from rotkehlchen.user_messages import MessagesAggregator
@@ -221,37 +222,7 @@ def test_binance_query_balances_unknown_asset(function_scope_binance):
     binance = function_scope_binance
 
     def mock_unknown_asset_return(url):  # pylint: disable=unused-argument
-        response = MockResponse(
-            200,
-            """
-            {
-  "makerCommission": 15,
-  "takerCommission": 15,
-  "buyerCommission": 0,
-  "sellerCommission": 0,
-  "canTrade": true,
-  "canWithdraw": true,
-  "canDeposit": true,
-  "updateTime": 123456789,
-  "balances": [
-    {
-      "asset": "BTC",
-      "free": "4723846.89208129",
-      "locked": "0.00000000"
-    }, {
-      "asset": "ETH",
-      "free": "4763368.68006011",
-      "locked": "0.00000000"
-    }, {
-      "asset": "IDONTEXIST",
-      "free": "5.0",
-      "locked": "0.0"
-    }, {
-      "asset": "ETF",
-      "free": "5.0",
-      "locked": "0.0"
-    }]}""")
-        return response
+        return MockResponse(200, BINANCE_BALANCES_RESPONSE)
 
     with patch.object(binance.session, 'get', side_effect=mock_unknown_asset_return):
         # Test that after querying the assets only ETH and BTC are there
diff --git a/rotkehlchen/tests/utils/constants.py b/rotkehlchen/tests/utils/constants.py
--- a/rotkehlchen/tests/utils/constants.py
+++ b/rotkehlchen/tests/utils/constants.py
@@ -14,3 +14,4 @@
 A_BNB = Asset('BNB')
 A_USDT = Asset('USDT')
 A_LTC = Asset('LTC')
+A_CNY = Asset('CNY')
diff --git a/rotkehlchen/tests/utils/exchanges.py b/rotkehlchen/tests/utils/exchanges.py
--- a/rotkehlchen/tests/utils/exchanges.py
+++ b/rotkehlchen/tests/utils/exchanges.py
@@ -84,3 +84,33 @@
         }},
     "count": 5
 }"""
+
+
+BINANCE_BALANCES_RESPONSE = """
+{
+  "makerCommission": 15,
+  "takerCommission": 15,
+  "buyerCommission": 0,
+  "sellerCommission": 0,
+  "canTrade": true,
+  "canWithdraw": true,
+  "canDeposit": true,
+  "updateTime": 123456789,
+  "balances": [
+    {
+      "asset": "BTC",
+      "free": "4723846.89208129",
+      "locked": "0.00000000"
+    }, {
+      "asset": "ETH",
+      "free": "4763368.68006011",
+      "locked": "0.00000000"
+    }, {
+      "asset": "IDONTEXIST",
+      "free": "5.0",
+      "locked": "0.0"
+    }, {
+      "asset": "ETF",
+      "free": "5.0",
+      "locked": "0.0"
+}]}"""