refactor: hide consensus and fork checks in functions

src/modules/accounting/accounting.py

@@ -210,28 +210,26 @@ def get_updated_modules_stats(
         return list(module_stats.keys()), list(module_stats.values())
 
     @lru_cache(maxsize=1)
-    def _get_consensus_lido_state_pre_electra(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]:
-        lido_validators = self.w3.lido_validators.get_lido_validators(blockstamp)
-
-        count = len(lido_validators)
-        total_balance = Gwei(sum(int(validator.balance) for validator in lido_validators))
+    def _get_consensus_lido_state(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]:
 
-        logger.info({'msg': 'Calculate lido state on CL. (Validators count, Total balance in gwei)', 'value': (count, total_balance)})
-        return ValidatorsCount(count), ValidatorsBalance(total_balance)
-
-    @lru_cache(maxsize=1)
-    def _get_consensus_lido_state_post_electra(self, blockstamp: ReferenceBlockStamp) -> tuple[ValidatorsCount, ValidatorsBalance]:
         lido_validators = self.w3.lido_validators.get_lido_validators(blockstamp)
+        logger.info({'msg': 'Calculate Lido validators count', 'value': len(lido_validators)})
 
-        validators_count = len(lido_validators)
-        active_balance = sum(int(validator.balance) for validator in lido_validators)
-        state = self.w3.cc.get_state_view(blockstamp)
-        pending_deposits = self.w3.lido_validators.calculate_pending_deposits_sum(lido_validators, state.pending_deposits)
-        total_balance = Gwei(active_balance + pending_deposits)
+        total_balance = sum(int(validator.balance) for validator in lido_validators)
+        logger.info({'msg': 'Calculate Lido active balance (in Gwei)', 'value': total_balance})
+
+        consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash)
+        if consensus_version > 2:
+            spec = self.w3.cc.get_config_spec()
+            if blockstamp.ref_epoch >= int(spec.ELECTRA_FORK_EPOCH):
+                state = self.w3.cc.get_state_view(blockstamp)
+                pending_deposits = self.w3.lido_validators.calculate_pending_deposits_sum(
+                    lido_validators, state.pending_deposits
+                )
+                logger.info({'msg': 'Calculate Lido pending deposits (in Gwei)', 'value': pending_deposits})
+                total_balance += pending_deposits
 
-        logger.info(
-            {'msg': f'Calculate Lido state on CL. {validators_count=}, {active_balance=}, {pending_deposits=}, {total_balance=} (Gwei)'})
-        return ValidatorsCount(validators_count), ValidatorsBalance(total_balance)
+        return ValidatorsCount(len(lido_validators)), ValidatorsBalance(Gwei(total_balance))
 
     def _get_finalization_data(self, blockstamp: ReferenceBlockStamp) -> tuple[FinalizationShareRate, FinalizationBatches]:
         simulation = self.simulate_full_rebase(blockstamp)
@@ -278,15 +276,7 @@ def simulate_rebase_after_report(
         """
         To calculate how much withdrawal request protocol can finalize - needs finalization share rate after this report
         """
-        consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash)
-        if consensus_version in (1, 2):
-            validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp)
-        else:
-            spec = self.w3.cc.get_config_spec()
-            if blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH):
-                validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp)
-            else:
-                validators_count, cl_balance = self._get_consensus_lido_state_post_electra(blockstamp)
+        validators_count, cl_balance = self._get_consensus_lido_state(blockstamp)
 
         chain_conf = self.get_chain_config(blockstamp)
 
@@ -399,15 +389,7 @@ def _calculate_report_latest_version(self, consensus_version: int, blockstamp: R
 
     # fetches validators_count, cl_balance, withdrawal_balance, el_vault_balance, shares_to_burn
     def _calculate_rebase_report(self, blockstamp: ReferenceBlockStamp) -> RebaseReport:
-        consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash)
-        if consensus_version in (1, 2):
-            validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp)
-        else:
-            spec = self.w3.cc.get_config_spec()
-            if blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH):
-                validators_count, cl_balance = self._get_consensus_lido_state_pre_electra(blockstamp)
-            else:
-                validators_count, cl_balance = self._get_consensus_lido_state_post_electra(blockstamp)
+        validators_count, cl_balance = self._get_consensus_lido_state(blockstamp)
         withdrawal_vault_balance = self.w3.lido_contracts.get_withdrawal_balance(blockstamp)
         el_rewards_vault_balance = self.w3.lido_contracts.get_el_vault_balance(blockstamp)
         shares_requested_to_burn = self.get_shares_to_burn(blockstamp)

src/services/bunker_cases/abnormal_cl_rebase.py

@@ -182,35 +182,13 @@ def _calculate_cl_rebase_between_blocks(
             self.w3.cc.get_validators_no_cache(prev_blockstamp),
         )
 
-        # Get Lido validators' balances with WithdrawalVault balance
-        consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(ref_blockstamp.block_hash)
-
-        def _pre_electra_fork():
-            ref = self._get_lido_validators_balance_with_vault_pre_electra(
-                ref_blockstamp, self.lido_validators
-            )
-            prev = self._get_lido_validators_balance_with_vault_pre_electra(
-                prev_blockstamp, prev_lido_validators
-            )
-            return ref, prev
-
-        def _post_electra_fork():
-            ref = self._get_lido_validators_balance_with_vault_post_electra(
-                ref_blockstamp, self.lido_validators
-            )
-            prev = self._get_lido_validators_balance_with_vault_post_electra(
-                prev_blockstamp, prev_lido_validators
-            )
-            return ref, prev
+        ref_lido_balance_with_vault = self._get_lido_validators_balance_with_vault(
+            ref_blockstamp, self.lido_validators
+        )
 
-        if consensus_version in (1, 2):
-            ref_lido_balance_with_vault, prev_lido_balance_with_vault = _pre_electra_fork()
-        else:
-            spec = self.w3.cc.get_config_spec()
-            if ref_blockstamp.ref_epoch < int(spec.ELECTRA_FORK_EPOCH):
-                ref_lido_balance_with_vault, prev_lido_balance_with_vault = _pre_electra_fork()
-            else:
-                ref_lido_balance_with_vault, prev_lido_balance_with_vault = _post_electra_fork()
+        prev_lido_balance_with_vault = self._get_lido_validators_balance_with_vault(
+            prev_blockstamp, prev_lido_validators
+        )
 
         # Raw CL rebase is calculated as difference between reference and previous Lido validators' balances
         # Without accounting withdrawals from WithdrawalVault
@@ -234,31 +212,31 @@ def _post_electra_fork():
 
         return cl_rebase
 
-    def _get_lido_validators_balance_with_vault_pre_electra(
+    def _get_lido_validators_balance_with_vault(
         self, blockstamp: BlockStamp, lido_validators: list[LidoValidator]
     ) -> Gwei:
         """
         Get Lido validator balance with withdrawals vault balance
         """
-        real_cl_balance = AbnormalClRebase.calculate_validators_balance_sum(lido_validators)
-        withdrawals_vault_balance = int(
-            self.w3.from_wei(self.w3.lido_contracts.get_withdrawal_balance_no_cache(blockstamp), "gwei")
-        )
-        return Gwei(real_cl_balance + withdrawals_vault_balance)
+        consensus_version = self.w3.lido_contracts.accounting_oracle.get_consensus_version(blockstamp.block_hash)
 
-    def _get_lido_validators_balance_with_vault_post_electra(
-        self, blockstamp: BlockStamp, lido_validators: list[LidoValidator]
-    ) -> Gwei:
-        """
-        Get Lido validator balance with withdrawals vault balance
-        """
         real_cl_balance = AbnormalClRebase.calculate_validators_balance_sum(lido_validators)
-        state = self.w3.cc.get_state_view(blockstamp)
-        pending_deposits_sum = LidoValidatorsProvider.calculate_pending_deposits_sum(lido_validators, state.pending_deposits)
         withdrawals_vault_balance = int(
             self.w3.from_wei(self.w3.lido_contracts.get_withdrawal_balance_no_cache(blockstamp), "gwei")
         )
-        return Gwei(real_cl_balance + pending_deposits_sum + withdrawals_vault_balance)
+        total_balance = real_cl_balance + withdrawals_vault_balance
+
+        if consensus_version > 2:
+            epoch = EpochNumber(blockstamp.slot_number // self.c_conf.slots_per_epoch)
+            spec = self.w3.cc.get_config_spec()
+            if epoch >= int(spec.ELECTRA_FORK_EPOCH):
+                state = self.w3.cc.get_state_view(blockstamp)
+                pending_deposits_sum = LidoValidatorsProvider.calculate_pending_deposits_sum(
+                    lido_validators, state.pending_deposits
+                )
+                total_balance += pending_deposits_sum
+
+        return Gwei(total_balance)
 
     def _get_withdrawn_from_vault_between_blocks(
         self, prev_blockstamp: BlockStamp, ref_blockstamp: ReferenceBlockStamp