wip: checkpoints refactoring + tests

src/modules/csm/csm.py

@@ -7,14 +7,14 @@
 
 from src.metrics.prometheus.business import CONTRACT_ON_PAUSE
 from src.metrics.prometheus.duration_meter import duration_meter
-from src.modules.csm.checkpoint import CheckpointsFactory
-from src.modules.csm.state import InvalidState, State
+from src.modules.csm.checkpoint import CheckpointsIterator, CheckpointProcessor
+from src.modules.csm.state import State, InvalidState
 from src.modules.csm.tree import Tree
 from src.modules.csm.types import ReportData
 from src.modules.submodules.consensus import ConsensusModule
 from src.modules.submodules.oracle_module import BaseModule, ModuleExecuteDelay
 from src.modules.submodules.types import ZERO_HASH
-from src.providers.execution.contracts import CSFeeOracle
+from src.providers.execution.contracts.cs_fee_oracle import CSFeeOracle
 from src.types import BlockStamp, EpochNumber, ReferenceBlockStamp, SlotNumber, ValidatorIndex
 from src.utils.cache import global_lru_cache as lru_cache
 from src.utils.slot import get_first_non_missed_slot
@@ -222,24 +222,33 @@ def collect_data(self, blockstamp: BlockStamp) -> bool:
         self.state.validate_for_collect(l_epoch, r_epoch)
         self.state.status()
 
-        factory = CheckpointsFactory(self.w3.cc, converter, self.state)
-        checkpoints = factory.prepare_checkpoints(l_epoch, r_epoch, finalized_epoch)
+        if done := self.state.is_fulfilled:
+            logger.info({"msg": "All epochs are already processed. Nothing to collect"})
+            return done
 
+        checkpoints = CheckpointsIterator(
+            converter, min(self.state.unprocessed_epochs) or l_epoch, r_epoch, finalized_epoch
+        )
+        processor = CheckpointProcessor(self.w3.cc, self.state, converter, blockstamp)
+
+        new_processed_epochs = 0
         start = time.time()
         for checkpoint in checkpoints:
+
             if self.current_frame_range(self._receive_last_finalized_slot()) != (l_ref_slot, r_ref_slot):
-                logger.info({"msg": "Checkpoints were prepared for an outdated frame, stop proccessing"})
+                logger.info({"msg": "Checkpoints were prepared for an outdated frame, stop processing"})
                 raise ValueError("Outdated checkpoint")
 
-            if converter.get_epoch_by_slot(checkpoint.slot) > finalized_epoch:
+            if checkpoint.slot > blockstamp.slot_number:
                 logger.info({"msg": f"Checkpoint for slot {checkpoint.slot} is not finalized yet"})
                 break
 
-            logger.info({"msg": f"Processing checkpoint for slot {checkpoint.slot}"})
-            logger.info({"msg": f"Processing {len(checkpoint.duty_epochs)} epochs"})
-            checkpoint.process(blockstamp)
-        if checkpoints:
-            logger.info({"msg": f"All epochs were processed in {time.time() - start:.2f} seconds"})
+            new_processed_epochs += processor.exec(checkpoint)
+
+        if new_processed_epochs:
+            logger.info(
+                {"msg": f"Collecting data for {new_processed_epochs} epochs was done in {time.time() - start:.2f} sec"}
+            )
 
         return self.state.is_fulfilled
 

src/web3py/extensions/csm.py

@@ -13,7 +13,9 @@
 
 from src import variables
 from src.metrics.prometheus.business import FRAME_PREV_REPORT_REF_SLOT
-from src.providers.execution.contracts import CSFeeDistributor, CSFeeOracle, CSModule
+from src.providers.execution.contracts.cs_fee_distributor import CSFeeDistributor
+from src.providers.execution.contracts.cs_fee_oracle import CSFeeOracle
+from src.providers.execution.contracts.cs_module import CSModule
 from src.providers.ipfs import CIDv0, CIDv1, is_cid_v0
 from src.types import BlockStamp, SlotNumber
 from src.web3py.extensions.lido_validators import NodeOperatorId

tests/modules/csm/test_checkpoint.py

@@ -0,0 +1,61 @@
+import pytest
+
+from src.modules.csm.checkpoint import CheckpointsIterator, Checkpoint
+from src.modules.submodules.types import ChainConfig, FrameConfig
+from src.utils.web3converter import Web3Converter
+from tests.factory.configs import FrameConfigFactory, ChainConfigFactory
+
+
+@pytest.fixture
+def frame_config() -> FrameConfig:
+    return FrameConfigFactory.build(
+        epochs_per_frame=2,
+    )
+
+
+@pytest.fixture
+def chain_config() -> ChainConfig:
+    return ChainConfigFactory.build(
+        slots_per_epoch=12,
+        seconds_per_slot=3,
+        genesis_time=0,
+    )
+
+
+@pytest.fixture
+def converter(frame_config: FrameConfig, chain_config: ChainConfig) -> Web3Converter:
+    return Web3Converter(chain_config, frame_config)
+
+
+def test_checkpoints_iterator_min_epoch_is_not_reached(converter):
+    with pytest.raises(ValueError, match='Minimum checkpoint step is not reached yet'):
+        CheckpointsIterator(converter, 100, 600, 109)
+
+
+def test_checkpoints_iterator_r_epoch_is_changed_by_finalized(converter):
+    l_epoch = 100
+    r_epoch = 600
+    finalized_epoch = 110
+    expected = finalized_epoch - 1
+    iterator = CheckpointsIterator(converter, l_epoch, r_epoch, finalized_epoch)
+    assert r_epoch != iterator.r_epoch, "Right border should be changed"
+    assert iterator.r_epoch == expected, "Right border should be equal to the finalized epoch minus one"
+
+
+@pytest.mark.parametrize(
+    "l_epoch,r_epoch,finalized_epoch,expected_checkpoints",
+    [
+        (0, 255, 255, [Checkpoint(3071, list(range(0, 255)))]),
+        (15, 255, 255, [Checkpoint(3071, list(range(15, 255)))]),
+        (15, 255, 25, [Checkpoint(311, list(range(15, 25)))]),
+        (0, 255 * 2, 255 * 2, [Checkpoint(3071, list(range(0, 255))), Checkpoint(6131, list(range(255, 510)))]),
+        (15, 255 * 2, 350, [Checkpoint(3251, list(range(15, 270))), Checkpoint(4211, list(range(270, 350)))]),
+    ],
+)
+def test_checkpoints_iterator_given_checkpoints(converter, l_epoch, r_epoch, finalized_epoch, expected_checkpoints):
+    iterator = CheckpointsIterator(converter, l_epoch, r_epoch, finalized_epoch)
+    assert iter(iterator).checkpoints == expected_checkpoints
+
+
+def test_checkpoints_processor():
+    assert True