Move generation to the model

lidofinance · Sep 18, 2024 · 919f8da · 919f8da
1 parent 87ee53f
commit 919f8da
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/tests/factory/no_registry.py b/tests/factory/no_registry.py
@@ -1,5 +1,6 @@
 import random
 from itertools import count
+from typing import Any
 
 from faker import Faker
 from pydantic_factories import Use
@@ -10,7 +11,6 @@
 from tests.factory.web3_factory import Web3Factory
 from src.web3py.extensions.lido_validators import StakingModule, LidoValidator, NodeOperator
 
-
 faker = Faker()
 
 
@@ -47,6 +47,12 @@ class LidoValidatorFactory(Web3Factory):
     index: str = Use(lambda x: str(next(x)), count(1))
     balance: str = Use(lambda x: str(x), random.randrange(1, 10**9))
 
+    @classmethod
+    def build_with_activation_epoch_bound(cls, max_value: int, **kwargs: Any):
+        return cls.build(
+            validator=ValidatorStateFactory.build(activation_epoch=str(faker.pyint(max_value=max_value - 1))), **kwargs
+        )
+
 
 class NodeOperatorFactory(Web3Factory):
     __model__ = NodeOperator

diff --git a/tests/modules/ejector/test_iterator_v2.py b/tests/modules/ejector/test_iterator_v2.py
@@ -130,26 +130,21 @@ def test_eject_validator(iterator):
         ]
     )
 
-    def generate_validator_state_with_activation_epoch_bound():
-        return ValidatorStateFactory.build(
-            activation_epoch=str(faker.pyint(max_value=iterator.blockstamp.ref_epoch - 1))
-        )
-
     iterator.w3.lido_validators.get_lido_validators_by_node_operators = Mock(
         return_value={
             (1, 1): [
-                LidoValidatorFactory.build(validator=generate_validator_state_with_activation_epoch_bound()),
-                LidoValidatorFactory.build(validator=generate_validator_state_with_activation_epoch_bound()),
-                LidoValidatorFactory.build(validator=generate_validator_state_with_activation_epoch_bound()),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch),
             ],
             (1, 2): [
-                LidoValidatorFactory.build(validator=generate_validator_state_with_activation_epoch_bound()),
-                LidoValidatorFactory.build(validator=generate_validator_state_with_activation_epoch_bound()),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch),
             ],
             (2, 1): [
-                LidoValidatorFactory.build(index='8', validator=generate_validator_state_with_activation_epoch_bound()),
-                LidoValidatorFactory.build(index='7', validator=generate_validator_state_with_activation_epoch_bound()),
-                LidoValidatorFactory.build(index='6', validator=generate_validator_state_with_activation_epoch_bound()),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch, index='8'),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch, index='7'),
+                LidoValidatorFactory.build_with_activation_epoch_bound(iterator.blockstamp.ref_epoch, index='6'),
             ],
         }
     )