Repository: rotki/rotki

Issue:
Finish the moving of eth2 code into a module
## Abstract

With https://github.com/rotki/rotki/pull/2852 we moved eth2 code into a module.

But I think the module still does not implement the account addition/delete and init code. And also I think that all addresses are queried instead of those belonging in the module.

## Task

Implement the remaining part needed for ETH2 module


Discussion:

Resolving patch:
diff --git a/rotkehlchen/chain/ethereum/modules/eth2.py b/rotkehlchen/chain/ethereum/modules/eth2.py
--- a/rotkehlchen/chain/ethereum/modules/eth2.py
+++ b/rotkehlchen/chain/ethereum/modules/eth2.py
@@ -5,7 +5,7 @@
 import gevent
 
 log = logging.getLogger(__name__)
-from rotkehlchen.accounting.structures import Balance
+from rotkehlchen.accounting.structures import AssetBalance, Balance
 from rotkehlchen.chain.ethereum.eth2_utils import scrape_validator_daily_stats
 from rotkehlchen.chain.ethereum.typing import (
     DEPOSITING_VALIDATOR_PERFORMANCE,
@@ -14,10 +14,11 @@
     ValidatorDetails,
 )
 from rotkehlchen.chain.ethereum.utils import decode_event_data
-from rotkehlchen.constants.assets import A_ETH
+from rotkehlchen.constants.assets import A_ETH, A_ETH2
 from rotkehlchen.constants.ethereum import EthereumConstants
 from rotkehlchen.constants.timing import DAY_IN_SECONDS
 from rotkehlchen.db.eth2 import ETH2_DEPOSITS_PREFIX, DBEth2
+from rotkehlchen.errors import RemoteError
 from rotkehlchen.fval import FVal
 from rotkehlchen.history.price import query_usd_price_zero_if_error
 from rotkehlchen.inquirer import Inquirer
@@ -28,7 +29,6 @@
 from rotkehlchen.utils.misc import from_gwei, ts_now
 
 if TYPE_CHECKING:
-    from rotkehlchen.accounting.structures import AssetBalance
     from rotkehlchen.chain.ethereum.manager import EthereumManager
     from rotkehlchen.db.dbhandler import DBHandler
     from rotkehlchen.externalapis.beaconchain import BeaconChain
@@ -48,16 +48,18 @@ class Eth2(EthereumModule):
     """Module representation for Eth2"""
 
     def __init__(
-        self,
-        ethereum_manager: 'EthereumManager',
-        database: 'DBHandler',
-        premium: Optional[Premium],
-        msg_aggregator: MessagesAggregator,
+            self,
+            ethereum_manager: 'EthereumManager',
+            database: 'DBHandler',
+            premium: Optional[Premium],
+            msg_aggregator: MessagesAggregator,
+            beaconchain: 'BeaconChain',
     ) -> None:
         self.database = database
         self.premium = premium
         self.ethereum = ethereum_manager
         self.msg_aggregator = msg_aggregator
+        self.beaconchain = beaconchain
         self.last_stats_query_ts = 0
         self.validator_stats_queried = 0
 
@@ -205,9 +207,8 @@ def get_staking_deposits(
         deposits.sort(key=lambda deposit: (deposit.timestamp, deposit.log_index))
         return deposits
 
-    @staticmethod
     def get_balances(
-            beaconchain: 'BeaconChain',
+            self,
             addresses: List[ChecksumEthAddress],
     ) -> Dict[ChecksumEthAddress, Balance]:
         """May Raise RemoteError from beaconcha.in api"""
@@ -218,7 +219,7 @@ def get_balances(
         balance_mapping: Dict[ChecksumEthAddress, Balance] = defaultdict(Balance)
         # Map eth1 addresses to validators
         for address in addresses:
-            validators = beaconchain.get_eth1_address_validators(address)
+            validators = self.beaconchain.get_eth1_address_validators(address)
             if len(validators) == 0:
                 continue
 
@@ -235,7 +236,7 @@ def get_balances(
                     )
 
         # Get current balance of all validator indices
-        performance = beaconchain.get_performance(validator_indices)
+        performance = self.beaconchain.get_performance(validator_indices)
         for validator_index, entry in performance.items():
             amount = from_gwei(entry.balance)
             balance_mapping[index_to_address[validator_index]] += Balance(amount, amount * usd_price)  # noqa: E501
@@ -252,7 +253,6 @@ def get_balances(
 
     def get_details(
             self,
-            beaconchain: 'BeaconChain',
             addresses: List[ChecksumEthAddress],
     ) -> List[ValidatorDetails]:
         """Go through the list of eth1 addresses and find all eth2 validators associated
@@ -263,10 +263,10 @@ def get_details(
         index_to_address = {}
         index_to_pubkey = {}
         result = []
-        assert beaconchain.db is not None, 'Beaconchain db should be populated'
+        assert self.beaconchain.db is not None, 'Beaconchain db should be populated'
 
         for address in addresses:
-            validators = beaconchain.get_eth1_address_validators(address)
+            validators = self.beaconchain.get_eth1_address_validators(address)
             for validator in validators:
                 if validator.validator_index is None:
                     # for validators that are so early in the depositing queue that no
@@ -285,11 +285,11 @@ def get_details(
                 indices.append(validator.validator_index)
 
         # Get current balance of all validator indices
-        performance_result = beaconchain.get_performance(list(indices))
+        performance_result = self.beaconchain.get_performance(list(indices))
         for validator_index, entry in performance_result.items():
             stats = self.get_validator_daily_stats(
                 validator_index=validator_index,
-                msg_aggregator=beaconchain.msg_aggregator,
+                msg_aggregator=self.msg_aggregator,
             )
             result.append(ValidatorDetails(
                 validator_index=validator_index,
@@ -304,7 +304,7 @@ def get_details(
         for index in depositing_indices:
             stats = self.get_validator_daily_stats(
                 validator_index=index,
-                msg_aggregator=beaconchain.msg_aggregator,
+                msg_aggregator=self.msg_aggregator,
             )
             result.append(ValidatorDetails(
                 validator_index=index,
@@ -375,7 +375,19 @@ def on_startup(self) -> None:
         pass
 
     def on_account_addition(self, address: ChecksumEthAddress) -> Optional[List['AssetBalance']]:
-        pass
+        try:
+            result = self.get_balances([address])
+        except RemoteError as e:
+            self.msg_aggregator.add_error(
+                f'Did not manage to query beaconcha.in api for address {address} due to {str(e)}.'
+                f' If you have Eth2 staked balances the final balance results may not be accurate',
+            )
+            return None
+        balance = result.get(address)
+        if balance is None:
+            return None
+
+        return [AssetBalance(asset=A_ETH2, balance=balance)]
 
     def on_account_removal(self, address: ChecksumEthAddress) -> None:
         pass
diff --git a/rotkehlchen/chain/manager.py b/rotkehlchen/chain/manager.py
--- a/rotkehlchen/chain/manager.py
+++ b/rotkehlchen/chain/manager.py
@@ -348,6 +348,9 @@ def activate_module(self, module_name: ModuleName) -> Optional[EthereumModule]:
             return module  # already activated
 
         logger.debug(f'Activating {module_name} module')
+        kwargs = {}
+        if module_name == 'eth2':
+            kwargs['beaconchain'] = self.beaconchain
         klass = _module_name_to_class(module_name)
         # TODO: figure out the type here: class EthereumModule not callable.
         try:
@@ -356,6 +359,7 @@ def activate_module(self, module_name: ModuleName) -> Optional[EthereumModule]:
                 database=self.database,
                 premium=self.premium,
                 msg_aggregator=self.msg_aggregator,
+                **kwargs,
             )
         except ModuleInitializationFailure as e:
             logger.error(f'Failed to activate {module_name} due to: {str(e)}')
@@ -644,8 +648,6 @@ def modify_eth_account(
             self.balances.eth[account] = BalanceSheet(
                 assets=defaultdict(Balance, {A_ETH: Balance(amount, usd_value)}),
             )
-            # Check if the new account has any staked eth2 deposits
-            self.account_for_staked_eth2_balances([account], at_addition=True)
         elif append_or_remove == 'remove':
             if account not in self.accounts.eth:
                 raise InputError('Tried to remove a non existing ETH account')
@@ -1140,10 +1142,7 @@ def _add_protocol_balances(self) -> None:
                 self.totals.assets[A_ADX] += balance
 
         # Count ETH staked in Eth2 beacon chain
-        self.account_for_staked_eth2_balances(
-            addresses=self.queried_addresses_for_module('eth2'),
-            at_addition=False,
-        )
+        self.account_for_staked_eth2_balances(addresses=self.queried_addresses_for_module('eth2'))
         # Finally count the balances detected in various protocols in defi balances
         self.add_defi_balances_to_token_and_totals()
 
@@ -1211,21 +1210,19 @@ def add_defi_balances_to_token_and_totals(self) -> None:
     def account_for_staked_eth2_balances(
             self,
             addresses: List[ChecksumEthAddress],
-            at_addition: bool = False,
     ) -> None:
         eth2 = self.get_module('eth2')
         if eth2 is None:
             return  # no eth2 module active -- do nothing
 
-        if not at_addition:
-            # Before querying the new balances, delete the ones in memory if any
-            self.totals.assets.pop(A_ETH2, None)
-            for _, entry in self.balances.eth.items():
-                if A_ETH2 in entry.assets:
-                    del entry.assets[A_ETH2]
+        # Before querying the new balances, delete the ones in memory if any
+        self.totals.assets.pop(A_ETH2, None)
+        for _, entry in self.balances.eth.items():
+            if A_ETH2 in entry.assets:
+                del entry.assets[A_ETH2]
 
         try:
-            mapping = eth2.get_balances(self.beaconchain, addresses)
+            mapping = eth2.get_balances(addresses)
         except RemoteError as e:
             self.msg_aggregator.add_error(
                 f'Did not manage to query beaconcha.in api for addresses due to {str(e)}.'
@@ -1264,10 +1261,7 @@ def get_eth2_staking_details(self) -> List['ValidatorDetails']:
         eth2 = self.get_module('eth2')
         if eth2 is None:
             raise ModuleInactive('Cant query eth2 staking details since eth2 module is not active')
-        return eth2.get_details(
-            beaconchain=self.beaconchain,
-            addresses=self.queried_addresses_for_module('eth2'),
-        )
+        return eth2.get_details(addresses=self.queried_addresses_for_module('eth2'))
 
     @protect_with_lock()
     @cache_response_timewise()
@@ -1287,10 +1281,7 @@ def get_eth2_history_events(
             return []  # no need to bother querying before beacon chain launch
 
         defi_events = []
-        eth2_details = eth2.get_details(
-            beaconchain=self.beaconchain,
-            addresses=self.queried_addresses_for_module('eth2'),
-        )
+        eth2_details = eth2.get_details(addresses=self.queried_addresses_for_module('eth2'))
         for entry in eth2_details:
             if entry.validator_index is None:
                 continue  # don't query stats for validators without an index yet


Test patch:
diff --git a/rotkehlchen/tests/fixtures/eth2.py b/rotkehlchen/tests/fixtures/eth2.py
--- a/rotkehlchen/tests/fixtures/eth2.py
+++ b/rotkehlchen/tests/fixtures/eth2.py
@@ -11,6 +11,7 @@ def fixture_eth2(
         messages_aggregator,
         start_with_valid_premium,
         rotki_premium_credentials,
+        beaconchain,
 ):
     premium = None
     if start_with_valid_premium:
@@ -20,4 +21,5 @@ def fixture_eth2(
         database=database,
         premium=premium,
         msg_aggregator=messages_aggregator,
+        beaconchain=beaconchain,
     )
diff --git a/rotkehlchen/tests/unit/test_eth2.py b/rotkehlchen/tests/unit/test_eth2.py
--- a/rotkehlchen/tests/unit/test_eth2.py
+++ b/rotkehlchen/tests/unit/test_eth2.py
@@ -424,8 +424,8 @@ def test_get_eth2_balances_validator_not_yet_active(beaconchain, inquirer, eth2)
 
     Test for: https://github.com/rotki/rotki/issues/1888
     """
-    with _create_beacon_mock(beaconchain):
-        mapping = eth2.get_balances(beaconchain, [ADDR1])
+    with _create_beacon_mock(eth2.beaconchain):
+        mapping = eth2.get_balances([ADDR1])
 
     assert len(mapping) == 1
     amount = FVal('64.143716247')
@@ -438,8 +438,8 @@ def test_get_eth2_details_validator_not_yet_active(beaconchain, inquirer, price_
 
     Test for: https://github.com/rotki/rotki/issues/1888
     """
-    with _create_beacon_mock(beaconchain):
-        details = eth2.get_details(beaconchain=beaconchain, addresses=[ADDR1])
+    with _create_beacon_mock(eth2.beaconchain):
+        details = eth2.get_details(addresses=[ADDR1])
 
     for idx in range(0, 2):
         # basic check about daily stats but then delete since this is not what this test checks