Timeout modules on CSM failure

src/blockchain/deposit_strategy/deposit_module_recommender.py

@@ -0,0 +1,63 @@
+from datetime import datetime, timedelta
+from typing import Callable
+
+from web3 import Web3
+
+_TIMEOUTS = {
+    1: timedelta(minutes=10),
+    2: timedelta(minutes=10),
+}
+
+
+class DepositModuleRecommender:
+    def __init__(self, w3: Web3):
+        self._w3 = w3
+        self._module_timeouts = dict()
+
+    def get_preferred_to_deposit_modules(self, whitelist_modules: list[int]) -> list[int]:
+        module_ids = self._w3.lido.staking_router.get_staking_module_ids()
+        modules = self._w3.lido.staking_router.get_staking_module_digests(module_ids)
+
+        depositable_modules = list(filter(self._get_module_depositable_filter(whitelist_modules), modules))
+        modules_ids = self.prioritize_modules(depositable_modules)
+        return modules_ids
+
+    def _get_module_depositable_filter(self, whitelist_modules: list[int]) -> Callable:
+        def is_module_depositable(module: list) -> bool:
+            module_id = module[2][0]
+
+            if self._is_timeout_passed(module_id):
+                self.reset_timeout(module_id)
+
+            return (
+                module_id not in self._module_timeouts
+                and module_id in whitelist_modules
+                and self._w3.lido.staking_router.is_staking_module_active(module_id)
+                and self._w3.lido.deposit_security_module.can_deposit(module_id)
+            )
+
+        return is_module_depositable
+
+    @staticmethod
+    def prioritize_modules(modules: list) -> list[int]:
+        modules = sorted(
+            modules,
+            #      totalDepositedValidators - totalExitedValidators
+            key=lambda module: module[3][1] - module[3][0],
+        )
+
+        #       module_ids
+        return [module[2][0] for module in modules]
+
+    def set_timeout(self, module_id: int):
+        self._module_timeouts[module_id] = datetime.now()
+
+    def reset_timeout(self, module_id: int):
+        return self._module_timeouts.pop(module_id, None)
+
+    def _is_timeout_passed(self, module_id: int) -> bool:
+        if module_id not in _TIMEOUTS or module_id not in self._module_timeouts:
+            return True
+        init = self._module_timeouts[module_id]
+        now = datetime.now()
+        return abs(now - init) >= _TIMEOUTS[module_id]

src/bots/depositor.py

@@ -5,9 +5,9 @@
 
 import variables
 from blockchain.deposit_strategy.base_deposit_strategy import CSMDepositStrategy, DefaultDepositStrategy
+from blockchain.deposit_strategy.deposit_module_recommender import DepositModuleRecommender
 from blockchain.deposit_strategy.deposit_transaction_sender import Sender
 from blockchain.deposit_strategy.gas_price_calculator import GasPriceCalculator
-from blockchain.deposit_strategy.prefered_module_to_deposit import get_preferred_to_deposit_modules
 from blockchain.deposit_strategy.strategy import DepositStrategy
 from blockchain.executor import Executor
 from blockchain.typings import Web3
@@ -34,14 +34,17 @@
 
 logger = logging.getLogger(__name__)
 
+CSM_MODULE = 3
+
 
 def run_depositor(w3):
     logger.info({'msg': 'Initialize Depositor bot.'})
     sender = Sender(w3)
     gas_price_calculator = GasPriceCalculator(w3)
     base_deposit_strategy = DefaultDepositStrategy(w3, gas_price_calculator)
     csm_strategy = CSMDepositStrategy(w3, gas_price_calculator)
-    depositor_bot = DepositorBot(w3, sender, base_deposit_strategy, csm_strategy)
+    recommender = DepositModuleRecommender(w3)
+    depositor_bot = DepositorBot(w3, sender, base_deposit_strategy, csm_strategy, recommender)
 
     e = Executor(
         w3,
@@ -66,11 +69,13 @@ def __init__(
         sender: Sender,
         base_deposit_strategy: DefaultDepositStrategy,
         csm_strategy: CSMDepositStrategy,
+        recommender: DepositModuleRecommender,
     ):
         self.w3 = w3
         self._sender = sender
         self._general_strategy = base_deposit_strategy
         self._csm_strategy = csm_strategy
+        self._recommender = recommender
 
         transports = []
 
@@ -108,21 +113,39 @@ def __init__(
     def execute(self, block: BlockData) -> bool:
         self._check_balance()
 
-        modules_id = get_preferred_to_deposit_modules(self.w3, variables.DEPOSIT_MODULES_WHITELIST)
+        module_ids = self._recommender.get_preferred_to_deposit_modules(variables.DEPOSIT_MODULES_WHITELIST)
 
-        if not modules_id:
+        if not module_ids:
             # Read messages in case if no depositable modules for metrics
             self.message_storage.receive_messages()
 
-        for module_id in modules_id:
+        for module_id in module_ids:
             logger.info({'msg': f'Do deposit to module with id: {module_id}.'})
             try:
-                self._deposit_to_module(module_id)
+                success = self._deposit_to_module(module_id)
+                if self._timeout_modules(success, module_id):
+                    break
             except ModuleNotSupportedError as error:
                 logger.warning({'msg': 'Module not supported exception.', 'error': str(error)})
 
         return True
 
+    def _timeout_modules(self, deposit_result: bool, module_id: int) -> bool:
+        if module_id != CSM_MODULE:
+            return False
+
+        if deposit_result:
+            # on successful deposit to CSM reset timeouts to all the other modules
+            for module in variables.DEPOSIT_MODULES_WHITELIST:
+                self._recommender.reset_timeout(module)
+            return False
+
+        # on unsuccessful deposit to CSM set timeout to other modules
+        for module in variables.DEPOSIT_MODULES_WHITELIST:
+            if module != CSM_MODULE:
+                self._recommender.set_timeout(module)
+        return True
+
     def _check_balance(self):
         eth_chain_id = self.w3.eth.chain_id
 
@@ -174,7 +197,7 @@ def _deposit_to_module(self, module_id: int) -> bool:
         return False
 
     def _select_strategy(self, module_id) -> DepositStrategy:
-        if module_id == 3:
+        if module_id == CSM_MODULE:
             return self._csm_strategy
         return self._general_strategy
 

tests/blockchain/deposit_strategy/test_deposit_module_recommender.py

@@ -0,0 +1,82 @@
+import unittest
+from datetime import datetime, timedelta
+from unittest.mock import MagicMock, patch
+
+from blockchain.deposit_strategy.deposit_module_recommender import DepositModuleRecommender
+
+
+class TestDepositModuleRecommender(unittest.TestCase):
+    def setUp(self):
+        # Create a mock Web3 instance with necessary Lido staking_router and deposit_security_module attributes
+        self.mock_w3 = MagicMock()
+        self.recommender = DepositModuleRecommender(w3=self.mock_w3)
+
+    @patch('blockchain.deposit_strategy.deposit_module_recommender.datetime')
+    def test_set_timeout(self, mock_datetime):
+        """Test setting a timeout for a module."""
+        mock_datetime.now.return_value = datetime(2023, 10, 1, 12, 0, 0)
+        module_id = 1
+
+        self.recommender.set_timeout(module_id)
+        self.assertIn(module_id, self.recommender._module_timeouts)
+        self.assertEqual(self.recommender._module_timeouts[module_id], mock_datetime.now())
+
+    def test_reset_timeout(self):
+        """Test resetting a timeout for a module."""
+        module_id = 1
+        self.recommender._module_timeouts[module_id] = datetime.now()
+        self.recommender.reset_timeout(module_id)
+        self.assertNotIn(module_id, self.recommender._module_timeouts)
+
+    @patch('blockchain.deposit_strategy.deposit_module_recommender.datetime')
+    def test_is_timeout_passed(self, mock_datetime):
+        """Test that _is_timeout_passed correctly identifies passed timeouts."""
+        mock_datetime.now.return_value = datetime(2023, 10, 1, 12, 10, 0)
+        module_id = 1
+        self.recommender._module_timeouts[module_id] = mock_datetime.now() - timedelta(minutes=15)
+
+        # Timeout of 10 minutes for module_id 1 should be exceeded
+        self.assertTrue(self.recommender._is_timeout_passed(module_id))
+
+    def test_prioritize_modules(self):
+        """Test prioritization of modules based on validator count."""
+        modules = [
+            ['ModuleA', 'info', [1], [5, 10]],  # Difference: 5
+            ['ModuleB', 'info', [2], [15, 20]],  # Difference: 5
+            ['ModuleC', 'info', [3], [1, 5]],  # Difference: 4
+        ]
+        prioritized = self.recommender.prioritize_modules(modules)
+        self.assertEqual(prioritized, [3, 1, 2])  # Module C has smallest difference, followed by A and B
+
+    def test_get_preferred_to_deposit_modules(self):
+        """Test retrieval of depositable modules based on whitelist and active status."""
+        whitelist_modules = [1, 2]
+        self.mock_w3.lido.staking_router.get_staking_module_ids.return_value = [1, 2, 3]
+        self.mock_w3.lido.staking_router.get_staking_module_digests.return_value = [
+            ['ModuleA', 'info', [1], [10, 5]],  # Module 1, whitelisted and active
+            ['ModuleB', 'info', [2], [20, 15]],  # Module 2, whitelisted but inactive
+            ['ModuleC', 'info', [3], [5, 1]],  # Module 3, not whitelisted
+        ]
+
+        # Set module activity and deposit capability
+        self.mock_w3.lido.staking_router.is_staking_module_active.side_effect = lambda module_id: module_id != 2
+        self.mock_w3.lido.deposit_security_module.can_deposit.side_effect = lambda module_id: module_id == 1
+
+        preferred_modules = self.recommender.get_preferred_to_deposit_modules(whitelist_modules)
+        self.assertEqual(preferred_modules, [1])
+
+    @patch('blockchain.deposit_strategy.deposit_module_recommender.datetime')
+    def test_depositable_filter_with_timeout(self, mock_datetime):
+        """Test the depositable filter with module timeout handling."""
+        mock_datetime.now.return_value = datetime(2023, 10, 1, 12, 0, 0)
+        whitelist_modules = [1]
+        module = ['ModuleA', 'info', [1], [10, 5]]
+
+        # Set up a timeout that has not yet expired
+        self.recommender.set_timeout(1)
+        self.recommender._module_timeouts[1] = mock_datetime.now() - timedelta(minutes=5)
+        self.assertFalse(self.recommender._get_module_depositable_filter(whitelist_modules)(module))
+
+        # Update time to simulate expiration of timeout
+        self.recommender._module_timeouts[1] = mock_datetime.now() - timedelta(minutes=15)
+        self.assertTrue(self.recommender._get_module_depositable_filter(whitelist_modules)(module))

tests/bots/test_depositor.py

@@ -1,8 +1,9 @@
-from unittest.mock import Mock
+import unittest
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 import variables
-from bots.depositor import DepositorBot
+from bots.depositor import CSM_MODULE, DepositorBot
 
 from tests.conftest import COUNCIL_ADDRESS_1, COUNCIL_ADDRESS_2, COUNCIL_PK_1, COUNCIL_PK_2, DSM_OWNER
 from tests.utils.protocol_utils import get_deposit_message
@@ -15,12 +16,13 @@ def depositor_bot(
     base_deposit_strategy,
     block_data,
     csm_strategy,
+    module_recommender,
 ):
     variables.MESSAGE_TRANSPORTS = ''
-    variables.DEPOSIT_MODULES_WHITELIST = [1, 2]
-    web3_lido_unit.lido.staking_router.get_staking_module_ids = Mock(return_value=[1, 2])
+    variables.DEPOSIT_MODULES_WHITELIST = [1, 2, 3]
+    web3_lido_unit.lido.staking_router.get_staking_module_ids = Mock(return_value=[1, 2, 3])
     web3_lido_unit.eth.get_block = Mock(return_value=block_data)
-    yield DepositorBot(web3_lido_unit, deposit_transaction_sender, base_deposit_strategy, csm_strategy)
+    yield DepositorBot(web3_lido_unit, deposit_transaction_sender, base_deposit_strategy, csm_strategy, module_recommender)
 
 
 @pytest.fixture
@@ -203,6 +205,88 @@ def test_get_quorum(depositor_bot, setup_deposit_message):
     assert deposit_messages[3] in quorum
 
 
+class TestDepositBot(unittest.TestCase):
+    @pytest.fixture(autouse=True)
+    def inject_fixtures(self, request):
+        """Inject pytest fixtures into unittest.TestCase."""
+        self.bot = request.getfixturevalue('depositor_bot')
+
+    def setUp(self):
+        # Mock dependencies and setup the instance
+        self.bot._check_balance = MagicMock()
+        self.bot._recommender = MagicMock()
+        self.bot.message_storage = MagicMock()
+        self.bot._deposit_to_module = MagicMock()
+
+    @pytest.mark.unit
+    @patch('bots.depositor.variables.DEPOSIT_MODULES_WHITELIST', [1, 2, CSM_MODULE])
+    def test_execute_no_depositable_modules(self):
+        """Test execute behavior when there are no depositable modules."""
+        self.bot._recommender.get_preferred_to_deposit_modules.return_value = []
+
+        result = self.bot.execute(block=MagicMock())
+        self.bot._check_balance.assert_called_once()
+        self.bot.message_storage.receive_messages.assert_called_once()
+        self.assertTrue(result)
+
+    @pytest.mark.unit
+    @patch('bots.depositor.variables.DEPOSIT_MODULES_WHITELIST', [1, CSM_MODULE])
+    def test_execute_deposit_to_modules(self):
+        """Test execute with available modules and successful deposits."""
+        self.bot._recommender.get_preferred_to_deposit_modules.return_value = [1, CSM_MODULE]
+        self.bot._deposit_to_module.side_effect = lambda module_id: module_id == CSM_MODULE  # Only CSM_MODULE succeeds
+
+        result = self.bot.execute(block=MagicMock())
+
+        # Check the correct sequence of method calls
+        self.bot._check_balance.assert_called_once()
+        self.bot._deposit_to_module.assert_any_call(1)
+        self.bot._deposit_to_module.assert_any_call(CSM_MODULE)
+        self.assertTrue(result)
+
+    @pytest.mark.unit
+    @patch('bots.depositor.variables.DEPOSIT_MODULES_WHITELIST', [1, CSM_MODULE])
+    def test_timeout_modules_successful_csm_deposit(self):
+        """Test _timeout_modules with a successful deposit to CSM_MODULE, expecting timeouts to reset."""
+        self.bot._recommender.reset_timeout = MagicMock()
+        self.bot._recommender.set_timeout = MagicMock()
+
+        result = self.bot._timeout_modules(deposit_result=True, module_id=CSM_MODULE)
+
+        # Reset timeout should be called for all modules in whitelist except CSM_MODULE
+        self.bot._recommender.reset_timeout.assert_any_call(1)
+        self.bot._recommender.set_timeout.assert_not_called()
+        self.assertFalse(result)
+
+    @pytest.mark.unit
+    @patch('bots.depositor.variables.DEPOSIT_MODULES_WHITELIST', [1, CSM_MODULE])
+    def test_timeout_modules_unsuccessful_csm_deposit(self):
+        """Test _timeout_modules with an unsuccessful deposit to CSM_MODULE, expecting timeouts to be set."""
+        self.bot._recommender.reset_timeout = MagicMock()
+        self.bot._recommender.set_timeout = MagicMock()
+
+        result = self.bot._timeout_modules(deposit_result=False, module_id=CSM_MODULE)
+
+        # Set timeout should be called for all modules in whitelist except CSM_MODULE
+        self.bot._recommender.set_timeout.assert_any_call(1)
+        self.bot._recommender.reset_timeout.assert_not_called()
+        self.assertTrue(result)
+
+    @pytest.mark.unit
+    @patch('bots.depositor.variables.DEPOSIT_MODULES_WHITELIST', [1, CSM_MODULE])
+    def test_timeout_modules_non_csm_module(self):
+        """Test _timeout_modules with a non-CSM_MODULE deposit, expecting no timeout changes."""
+        self.bot._recommender.reset_timeout = MagicMock()
+        self.bot._recommender.set_timeout = MagicMock()
+
+        result = self.bot._timeout_modules(deposit_result=True, module_id=1)
+
+        # No timeouts should be modified for non-CSM_MODULE deposits
+        self.bot._recommender.reset_timeout.assert_not_called()
+        self.bot._recommender.set_timeout.assert_not_called()
+        self.assertFalse(result)
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize(
     'web3_provider_integration,module_id',
@@ -218,9 +302,10 @@ def test_depositor_bot(
     csm_strategy_integration,
     module_id,
     add_accounts_to_guardian,
+    module_recommender_integration,
 ):
     # Define the whitelist of deposit modules
-    variables.DEPOSIT_MODULES_WHITELIST = [1, 2]
+    variables.DEPOSIT_MODULES_WHITELIST = [1, 2, 3]
 
     # Set the balance for the first account
     web3_lido_integration.provider.make_request(
@@ -265,6 +350,7 @@ def test_depositor_bot(
         deposit_transaction_sender_integration,
         base_deposit_strategy_integration,
         csm_strategy_integration,
+        module_recommender_integration,
     )
 
     # Clear the message storage and execute the bot without any messages

tests/fixtures/__init__.py

@@ -24,6 +24,8 @@
     deposit_transaction_sender_integration,
     gas_price_calculator,
     gas_price_calculator_integration,
+    module_recommender,
+    module_recommender_integration,
 )
 
 __all__ = [
@@ -48,4 +50,6 @@
     'deposit_transaction_sender_integration',
     'csm_strategy',
     'csm_strategy_integration',
+    'module_recommender',
+    'module_recommender_integration',
 ]