From 6ee40a828a0060cf75b5898723f3ab24be91b387 Mon Sep 17 00:00:00 2001 From: hweawer Date: Mon, 11 Nov 2024 08:39:53 +0100 Subject: [PATCH 1/4] timeout modules --- .../deposit_module_recommender.py | 63 +++++++++++++++++++ src/bots/depositor.py | 34 ++++++++-- 2 files changed, 91 insertions(+), 6 deletions(-) create mode 100644 src/blockchain/deposit_strategy/deposit_module_recommender.py diff --git a/src/blockchain/deposit_strategy/deposit_module_recommender.py b/src/blockchain/deposit_strategy/deposit_module_recommender.py new file mode 100644 index 0000000..4a0d8f2 --- /dev/null +++ b/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] diff --git a/src/bots/depositor.py b/src/bots/depositor.py index 5153b76..bc20bc2 100644 --- a/src/bots/depositor.py +++ b/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,6 +34,8 @@ logger = logging.getLogger(__name__) +CSM_MODULE = 3 + def run_depositor(w3): logger.info({'msg': 'Initialize Depositor bot.'}) @@ -66,11 +68,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,16 +112,18 @@ 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)}) @@ -174,10 +180,26 @@ 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 + 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_module_status(self, module_id: int) -> bool: """Returns True if module is ready for deposit""" ready = self.w3.lido.staking_router.is_staking_module_active(module_id) From 436eb3d7e5ee7ea4b08455a97250a228ba4c5201 Mon Sep 17 00:00:00 2001 From: hweawer Date: Mon, 11 Nov 2024 08:51:56 +0100 Subject: [PATCH 2/4] fix constructors in tests --- src/bots/depositor.py | 3 ++- tests/bots/test_depositor.py | 5 ++++- tests/fixtures/__init__.py | 4 ++++ tests/fixtures/strategy.py | 11 +++++++++++ 4 files changed, 21 insertions(+), 2 deletions(-) diff --git a/src/bots/depositor.py b/src/bots/depositor.py index bc20bc2..504b8c8 100644 --- a/src/bots/depositor.py +++ b/src/bots/depositor.py @@ -43,7 +43,8 @@ def run_depositor(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, diff --git a/tests/bots/test_depositor.py b/tests/bots/test_depositor.py index 2f62f14..61b52d1 100644 --- a/tests/bots/test_depositor.py +++ b/tests/bots/test_depositor.py @@ -15,12 +15,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]) 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 @@ -218,6 +219,7 @@ 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] @@ -265,6 +267,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 diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py index f5e7bda..a008203 100644 --- a/tests/fixtures/__init__.py +++ b/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', ] diff --git a/tests/fixtures/strategy.py b/tests/fixtures/strategy.py index 93c40ed..2d0bdec 100644 --- a/tests/fixtures/strategy.py +++ b/tests/fixtures/strategy.py @@ -1,5 +1,6 @@ import pytest 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 @@ -42,3 +43,13 @@ def csm_strategy(web3_lido_unit, gas_price_calculator): @pytest.fixture def csm_strategy_integration(web3_lido_integration, gas_price_calculator_integration): yield CSMDepositStrategy(web3_lido_integration, gas_price_calculator_integration) + + +@pytest.fixture +def module_recommender(web3_lido_unit): + yield DepositModuleRecommender(web3_lido_unit) + + +@pytest.fixture +def module_recommender_integration(web3_lido_integration): + yield DepositModuleRecommender(web3_lido_integration) From 4d73699d6760b4db8a8ceb2b043164d36313813d Mon Sep 17 00:00:00 2001 From: hweawer Date: Mon, 11 Nov 2024 12:42:49 +0100 Subject: [PATCH 3/4] Add csm module in tests --- tests/bots/test_depositor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/bots/test_depositor.py b/tests/bots/test_depositor.py index 61b52d1..30cf700 100644 --- a/tests/bots/test_depositor.py +++ b/tests/bots/test_depositor.py @@ -18,8 +18,8 @@ def depositor_bot( 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, module_recommender) From 278bd2315e8579328964281f107b9c8b606101a4 Mon Sep 17 00:00:00 2001 From: hweawer Date: Mon, 11 Nov 2024 13:26:10 +0100 Subject: [PATCH 4/4] Unit tests --- src/bots/depositor.py | 32 +++---- .../test_deposit_module_recommender.py | 82 +++++++++++++++++ tests/bots/test_depositor.py | 89 ++++++++++++++++++- 3 files changed, 184 insertions(+), 19 deletions(-) create mode 100644 tests/blockchain/deposit_strategy/test_deposit_module_recommender.py diff --git a/src/bots/depositor.py b/src/bots/depositor.py index 504b8c8..c5da0e1 100644 --- a/src/bots/depositor.py +++ b/src/bots/depositor.py @@ -130,6 +130,22 @@ def execute(self, block: BlockData) -> bool: 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 @@ -185,22 +201,6 @@ def _select_strategy(self, module_id) -> DepositStrategy: return self._csm_strategy return self._general_strategy - 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_module_status(self, module_id: int) -> bool: """Returns True if module is ready for deposit""" ready = self.w3.lido.staking_router.is_staking_module_active(module_id) diff --git a/tests/blockchain/deposit_strategy/test_deposit_module_recommender.py b/tests/blockchain/deposit_strategy/test_deposit_module_recommender.py new file mode 100644 index 0000000..fad13d8 --- /dev/null +++ b/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)) diff --git a/tests/bots/test_depositor.py b/tests/bots/test_depositor.py index 30cf700..80de8e5 100644 --- a/tests/bots/test_depositor.py +++ b/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 @@ -204,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', @@ -222,7 +305,7 @@ def test_depositor_bot( 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(