diff --git a/backend/consensus/vrf.py b/backend/consensus/vrf.py index 55c953600..6135e0d2d 100644 --- a/backend/consensus/vrf.py +++ b/backend/consensus/vrf.py @@ -6,18 +6,26 @@ def get_validators_for_transaction( nodes: list[dict], num_validators: int | None = None, - rng=np.random.default_rng(seed=int(datetime.now().timestamp())), + rng: np.random.Generator | None = None, ) -> list[dict]: """ Returns subset of validators for a transaction. The selelction and order is given by a random sampling based on the stake of the validators. """ + if rng is None: + rng = np.random.default_rng(seed=int(datetime.now().timestamp())) + if num_validators is None: num_validators = DEFAULT_VALIDATORS_COUNT num_validators = min(num_validators, len(nodes)) total_stake = sum(validator["stake"] for validator in nodes) + if total_stake == 0: + raise ValueError( + "Cannot select validators: all stakes are zero. " + "At least one validator must have a positive stake." + ) probabilities = [validator["stake"] / total_stake for validator in nodes] selected_validators = rng.choice( diff --git a/tests/unit/test_vrf.py b/tests/unit/test_vrf.py index 95473770e..4760b3d35 100644 --- a/tests/unit/test_vrf.py +++ b/tests/unit/test_vrf.py @@ -1,5 +1,6 @@ from backend.consensus.vrf import get_validators_for_transaction from unittest.mock import Mock +import pytest def list_of_dicts_to_set(list_of_dicts: list[dict]) -> set: @@ -66,3 +67,21 @@ def choice_mock(a, p, size, replace): rng.choice.assert_called_once() assert validators == [{"stake": 3}, {"stake": 2}, {"stake": 1}] + + +def test_get_validators_zero_stake_raises(): + """When all stakes are zero, selection should raise ValueError.""" + nodes = [{"stake": 0}, {"stake": 0}] + with pytest.raises(ValueError, match="all stakes are zero"): + get_validators_for_transaction(nodes, 1) + + +def test_get_validators_rng_not_shared_between_calls(): + """Each call without an explicit rng should get a fresh generator.""" + nodes = [{"stake": 1}, {"stake": 2}, {"stake": 3}] + # Call twice without passing rng — the default must be created fresh each time. + v1 = get_validators_for_transaction(nodes, 1) + v2 = get_validators_for_transaction(nodes, 1) + # Both should succeed (no shared mutable state corruption). + assert len(v1) == 1 + assert len(v2) == 1