Raise exception when the reservoir of values to draw from is less than required successes.

pull/3294/head
derekpierre 2023-11-07 14:08:31 -05:00 committed by David Núñez
parent b0f9fe1897
commit bad425f17f
1 changed files with 11 additions and 0 deletions

View File

@ -49,6 +49,9 @@ class MergedReservoir:
else:
return None
def __len__(self):
return len(self.values) + len(self.reservoir)
class PrefetchStrategy:
"""
@ -58,10 +61,18 @@ class PrefetchStrategy:
"""
def __init__(self, reservoir: MergedReservoir, need_successes: int):
if len(reservoir) < need_successes:
raise ValueError(
f"Insufficient staking providers ({len(reservoir.values)}) to draw {need_successes}"
)
self.reservoir = reservoir
self.need_successes = need_successes
def __call__(self, successes: int) -> Optional[List[ChecksumAddress]]:
if successes > self.need_successes:
# TODO: we could raise here, but we were returning `None` before
return None
batch = []
for i in range(self.need_successes - successes):
value = self.reservoir()