Add has_enough_shares method

Scille · Jul 29, 2023 · 873d143 · 873d143
1 parent 859ca39
commit 873d143
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/parsec/core/invite/claimer.py b/parsec/core/invite/claimer.py
@@ -217,6 +217,9 @@ def remaining_recipients(self) -> list[ShamirRecoveryRecipient]:
             recipient for recipient in self.recipients if recipient.user_id not in self.recovered
         ]
 
+    def has_enough_shares(self) -> bool:
+        return len(self.shares) >= self.threshold
+
     def get_initial_ctx(self, recipient: ShamirRecoveryRecipient) -> ShamirRecoveryClaimInitialCtx:
         return ShamirRecoveryClaimInitialCtx(
             recipient=recipient,
@@ -228,10 +231,9 @@ def get_initial_ctx(self, recipient: ShamirRecoveryRecipient) -> ShamirRecoveryC
 
     def add_shares(self, recipient_user_id: UserID, shares: list[ShamirShare]) -> bool:
         self.recovered[recipient_user_id] = shares
-        return len(self.shares) >= self.threshold
+        return self.has_enough_shares()
 
     async def retrieve_recovery_device(self) -> LocalDevice:
-        assert len(self.shares) >= self.threshold
         recovered_secret = ShamirRecoverySecret.load(
             shamir_recover_secret(self.threshold, self.shares)
         )

diff --git a/tests/core/test_shamir.py b/tests/core/test_shamir.py
@@ -330,6 +330,7 @@ async def _run_first_alice_claimer():
                         alice_ctx = await alice_ctx.do_wait_peer_trust()
                         enough_shares = await alice_ctx.do_recover_share()
                         assert not enough_shares
+                        assert not prelude_ctx.has_enough_shares()
                         assert len(prelude_ctx.shares) == 3
                         assert prelude_ctx.remaining_recipients == [adam_recipient]
                         alice_done_with_bob.set()
@@ -361,6 +362,7 @@ async def _run_first_alice_claimer():
                     alice_ctx = await alice_ctx.do_wait_peer_trust()
                     enough_shares = await alice_ctx.do_recover_share()
                     assert enough_shares
+                    assert prelude_ctx.has_enough_shares()
                     assert len(prelude_ctx.shares) == 5
                     assert prelude_ctx.remaining_recipients == []
                     await adam_done_with_alice.wait()