From bad425f17f211e2eb4d5e490cbba79dde5206379 Mon Sep 17 00:00:00 2001 From: derekpierre Date: Tue, 7 Nov 2023 14:08:31 -0500 Subject: [PATCH] Raise exception when the reservoir of values to draw from is less than required successes. --- nucypher/policy/reservoir.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/nucypher/policy/reservoir.py b/nucypher/policy/reservoir.py index 5a33ca831..fc9f14882 100644 --- a/nucypher/policy/reservoir.py +++ b/nucypher/policy/reservoir.py @@ -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()