Merge pull request #515 from lidofinance/fix-defaultdict-access

fix: defaultdict access dont raise KeyError
lidofinance · Sep 9, 2024 · 13366d5 · 13366d5
2 parents dbd4a62 + 339438c
commit 13366d5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/src/modules/csm/csm.py b/src/modules/csm/csm.py
@@ -218,18 +218,18 @@ def calculate_distribution(
                 continue
 
             for v in validators:
-                try:
-                    aggr = self.state.data[ValidatorIndex(int(v.index))]
-                except KeyError:
+                aggr = self.state.data.get(ValidatorIndex(int(v.index)))
+
+                if aggr is None:
                     # It's possible that the validator is not assigned to any duty, hence it's performance
                     # is not presented in the aggregates (e.g. exited, pending for activation etc).
                     continue
-                else:
-                    if v.validator.slashed is True:
-                        # It means that validator was active during the frame and got slashed and didn't meet the exit
-                        # epoch, so we should not count such validator for operator's share.
-                        log.operators[no_id].validators[v.index].slashed = True
-                        continue
+
+                if v.validator.slashed is True:
+                    # It means that validator was active during the frame and got slashed and didn't meet the exit
+                    # epoch, so we should not count such validator for operator's share.
+                    log.operators[no_id].validators[v.index].slashed = True
+                    continue
 
                 if aggr.perf > threshold:
                     # Count of assigned attestations used as a metrics of time

diff --git a/tests/modules/csm/test_csm_module.py b/tests/modules/csm/test_csm_module.py
@@ -148,7 +148,7 @@ def test_calculate_distribution(module: CSOracle, csm: CSM):
         NodeOperatorId(4),
         NodeOperatorId(5),
         NodeOperatorId(6),
-        NodeOperatorId(7),
+        # NodeOperatorId(7), # Missing in state
         NodeOperatorId(8),
         NodeOperatorId(9),
     )