fix: fork tests

Author: vgorkavenko

.github/workflows/ci-daily-checks.yml

@@ -22,7 +22,7 @@ jobs:
           CONSENSUS_CLIENT_URI: ${{ secrets.CONSENSUS_CLIENT_URI }}
           KEYS_API_URI: ${{ secrets.KEYS_API_URI }}
           LIDO_LOCATOR_ADDRESS: ${{ secrets.LIDO_LOCATOR_ADDRESS }}
-          STAKING_MODULE_ADDRESS: ${{ secrets.CSM_MODULE_ADDRESS }}
+          STAKING_MODULE_ADDRESS: ${{ secrets.STAKING_MODULE_ADDRESS }}
           PINATA_JWT: ${{ secrets.PINATA_JWT }}
           PINATA_DEDICATED_GATEWAY_URL: ${{ secrets.PINATA_DEDICATED_GATEWAY_URL }}
           PINATA_DEDICATED_GATEWAY_TOKEN: ${{ secrets.PINATA_DEDICATED_GATEWAY_TOKEN}}
@@ -39,7 +39,7 @@ jobs:
           CONSENSUS_CLIENT_URI=${CONSENSUS_CLIENT_URI}
           KEYS_API_URI=${KEYS_API_URI}
           LIDO_LOCATOR_ADDRESS=${LIDO_LOCATOR_ADDRESS}
-          CSM_MODULE_ADDRESS=${CSM_MODULE_ADDRESS}
+          STAKING_MODULE_ADDRESS=${STAKING_MODULE_ADDRESS}
           GW3_ACCESS_KEY=${GW3_ACCESS_KEY}
           GW3_SECRET_KEY=${GW3_SECRET_KEY}
           PINATA_JWT=${PINATA_JWT}

src/modules/oracles/staking_modules/base.py

@@ -177,11 +177,9 @@ def collect_data(self) -> bool:
     def build_report(self, blockstamp: ReferenceBlockStamp) -> tuple:
         self.validate_state(blockstamp)
 
-        last_report = self._get_last_report(blockstamp)
+        distribution, last_report = self.calculate_distribution(blockstamp)
         rewards_tree_root, rewards_cid = last_report.rewards_tree_root, last_report.rewards_tree_cid
 
-        distribution = self.calculate_distribution(blockstamp, last_report)
-
         if distribution.total_rewards:
             rewards_tree = self.make_rewards_tree(distribution.total_rewards_map)
             rewards_tree_root = rewards_tree.root
@@ -218,10 +216,11 @@ def _get_last_report(self, blockstamp: ReferenceBlockStamp) -> LastReport:
         return LastReport.load(self.w3, blockstamp, current_frame)
 
     @lru_cache(maxsize=1)
-    def calculate_distribution(self, blockstamp: ReferenceBlockStamp, last_report: LastReport) -> DistributionResult:
+    def calculate_distribution(self, blockstamp: ReferenceBlockStamp) -> tuple[DistributionResult, LastReport]:
+        last_report = self._get_last_report(blockstamp)
         distribution = Distribution(self.w3, self.converter(blockstamp), self.state)
         result = distribution.calculate(blockstamp, last_report)
-        return result
+        return result, last_report
 
     def is_main_data_submitted(self, blockstamp: BlockStamp) -> bool:
         last_ref_slot = self.w3.staking_module.get_last_processing_ref_slot(blockstamp)

tests/fork/test_csm_oracle_cycle.py

@@ -1,6 +1,7 @@
 from threading import Thread
 
 import pytest
+import uvicorn
 
 from unittest.mock import patch
 from pathlib import Path
@@ -12,11 +13,13 @@
 from src.modules.sidecars.performance.collector.collector import PerformanceCollector
 from src.modules.common.types import FrameConfig
 from src.modules.sidecars.performance.common.db import Duty
-from src.modules.sidecars.performance.web.server import serve
+from src.modules.sidecars.performance.web.server import app
 from src.utils.range import sequence
 from src.web3py.types import Web3
 from tests.fork.conftest import first_slot_of_epoch
 
+# pylint: disable=protected-access
+
 
 @pytest.fixture()
 def hash_consensus_bin():
@@ -40,12 +43,19 @@ def performance_local_db(testrun_path):
 
     def mock_get_database_url(self):
         db_path = Path(testrun_path) / "test_duties.db"
+        db_path.parent.mkdir(parents=True, exist_ok=True)
         return f"sqlite:///{db_path}"
 
+    def mock_build_engine(self, connect_timeout):
+        return create_engine(
+            self._get_database_url(),
+            echo=False,
+            pool_pre_ping=True,
+        )
+
     def mock_init(self, *args, **kwargs):
-        # pylint: disable=protected-access
-        self.engine = create_engine(self._get_database_url(), echo=False)
-        # pylint: disable=protected-access
+        self._statement_timeout_ms = kwargs.get('statement_timeout_ms')
+        self.engine = self._build_engine(kwargs.get('connect_timeout'))
         self._setup_database()
 
     table = Duty.__table__
@@ -54,8 +64,9 @@ def mock_init(self, *args, **kwargs):
             table.c[col_name].type = JSON()
 
     with patch('src.modules.sidecars.performance.common.db.DutiesDB._get_database_url', mock_get_database_url):
-        with patch('src.modules.sidecars.performance.common.db.DutiesDB.__init__', mock_init):
-            yield mock_get_database_url, mock_init
+        with patch('src.modules.sidecars.performance.common.db.DutiesDB._build_engine', mock_build_engine):
+            with patch('src.modules.sidecars.performance.common.db.DutiesDB.__init__', mock_init):
+                yield mock_get_database_url, mock_init
 
 
 @pytest.fixture()
@@ -65,7 +76,9 @@ def performance_collector(performance_local_db, web3: Web3, frame_config: FrameC
 
 @pytest.fixture()
 def performance_web_server(performance_local_db):
-    Thread(target=serve, daemon=True).start()
+    Thread(
+        target=uvicorn.run, args=(app,), kwargs={'host': '127.0.0.1', 'port': 9020, 'log_level': 'error'}, daemon=True
+    ).start()
     yield
 
 
@@ -98,19 +111,12 @@ def missed_initial_frame(frame_config: FrameConfig, cycle_iterations):
 @pytest.mark.fork
 @pytest.mark.parametrize(
     'module',
-    [
-        csm_module,
-        # cm_module
-    ],
+    [csm_module, cm_module],
     indirect=True,
 )
 @pytest.mark.parametrize(
     'running_finalized_slots',
-    [
-        # start_before_initial_epoch,
-        start_after_initial_epoch,
-        # missed_initial_frame
-    ],
+    [start_before_initial_epoch, start_after_initial_epoch, missed_initial_frame],
     indirect=True,
 )
 def test_csm_module_report(

tests/modules/csm/test_csm_module.py

@@ -756,9 +756,7 @@ class BuildReportTestParam:
 def test_build_report(module: CSPerformanceOracle, param: BuildReportTestParam):
     module.validate_state = Mock()
     module.report_contract.get_consensus_version = Mock(return_value=1)
-    module._get_last_report = Mock(return_value=param.last_report)
-    # mock current frame
-    module.calculate_distribution = param.curr_distribution
+    module.calculate_distribution = Mock(return_value=(param.curr_distribution(), param.last_report))
     module.make_rewards_tree = Mock(return_value=Mock(root=param.curr_rewards_tree_root))
     module.make_strikes_tree = Mock(return_value=Mock(root=param.curr_strikes_tree_root))
     module.publish_tree = Mock(
@@ -843,20 +841,23 @@ def test_calculate_distribution_lru_cache(module: CSPerformanceOracle):
 
         module.converter = Mock()
         module.state = Mock()
+        module._get_last_report = Mock(return_value=last_report)
 
-        result1 = module.calculate_distribution(blockstamp, last_report)
+        result1, last_report1 = module.calculate_distribution(blockstamp)
 
-        result2 = module.calculate_distribution(blockstamp, last_report)
+        result2, last_report2 = module.calculate_distribution(blockstamp)
 
         assert result1 is result2
+        assert last_report1 is last_report2
         assert result1 is mock_distribution_result
+        assert last_report1 is last_report
 
         assert MockDistribution.call_count == 1
         assert mock_distribution_instance.calculate.call_count == 1
 
         module.calculate_distribution.cache_clear()
 
-        result3 = module.calculate_distribution(blockstamp, last_report)
+        result3, last_report3 = module.calculate_distribution(blockstamp)
 
         assert MockDistribution.call_count == 2
         assert result3 is mock_distribution_result

tests/providers_clients/test_keys_api_client.py

@@ -80,11 +80,11 @@ def test_get_used_module_operators_keys__csm_module__response_data_is_valid(
         empty_blockstamp,
     ):
         csm_module_operators_keys = keys_api_client.get_used_module_operators_keys(
-            module_address=variables.CSM_MODULE_ADDRESS,  # type: ignore
+            module_address=variables.STAKING_MODULE_ADDRESS,  # type: ignore
             blockstamp=empty_blockstamp,
         )
 
-        assert csm_module_operators_keys['module']['stakingModuleAddress'] == variables.CSM_MODULE_ADDRESS
+        assert csm_module_operators_keys['module']['stakingModuleAddress'] == variables.STAKING_MODULE_ADDRESS
         assert csm_module_operators_keys['module']['id'] >= 0
         assert len(csm_module_operators_keys['keys']) > 0
         assert len(csm_module_operators_keys['operators']) > 0
@@ -97,7 +97,7 @@ def test_get_used_module_operators_keys__csm_module__response_data_is_valid(
         for operator in csm_module_operators_keys['operators']:
             assert operator['index'] >= 0
             assert Web3.is_address(operator['rewardAddress'])
-            assert operator['moduleAddress'] == variables.CSM_MODULE_ADDRESS
+            assert operator['moduleAddress'] == variables.STAKING_MODULE_ADDRESS
 
     def test_get_status__response_version_is_allowed(
         self,