diff --git a/src/_zkapauthorizer/tests/matchers.py b/src/_zkapauthorizer/tests/matchers.py index 2c90d8e74457d3a196de64dce3743db381570776..7d214e3882594b82cd30200721bbe3a73db0ce4b 100644 --- a/src/_zkapauthorizer/tests/matchers.py +++ b/src/_zkapauthorizer/tests/matchers.py @@ -135,15 +135,17 @@ def leases_current(relevant_storage_indexes, now, min_lease_remaining): servers for which the leases on the given storage indexes do not expire before ``min_lease_remaining``. """ + def get_relevant_stats(storage_server): + for (storage_index, shares) in storage_server.buckets.items(): + if storage_index in relevant_storage_indexes: + for (sharenum, stat) in shares.items(): + yield stat + + return AfterPreprocessing( # Get share stats for storage indexes we should have # visited and maintained. - lambda storage_server: list( - stat - for (storage_index, shares) in storage_server.buckets.items() - if storage_index in relevant_storage_indexes - for (sharenum, stat) in shares.items() - ), + lambda storage_server: list(get_relevant_stats(storage_server)), AllMatch( AfterPreprocessing( # Lease expiration for anything visited must be