From c5bf021c246d40172fd834a05d369e8892f3b9c6 Mon Sep 17 00:00:00 2001
From: derekpierre <derek.pierre@gmail.com>
Date: Mon, 15 Jan 2024 16:11:45 -0500
Subject: [PATCH] Update tests to have ferveo public keys obtained from
 CoordinatorAgent/Coordinator contract.

Co-authored-by: Kieran Prasch <kieranprasch@gmail.com>
---
 .../blockchain/test_integration_dkg_ritual.py         |  4 ++++
 tests/mock/coordinator.py                             |  8 ++++----
 tests/unit/test_ritualist.py                          | 11 +++++++++--
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/tests/integration/blockchain/test_integration_dkg_ritual.py b/tests/integration/blockchain/test_integration_dkg_ritual.py
index ba3d5df31..00e3da79a 100644
--- a/tests/integration/blockchain/test_integration_dkg_ritual.py
+++ b/tests/integration/blockchain/test_integration_dkg_ritual.py
@@ -12,6 +12,7 @@ from web3.datastructures import AttributeDict
 from nucypher.blockchain.eth.agents import CoordinatorAgent
 from nucypher.blockchain.eth.signers.software import Web3Signer
 from nucypher.characters.lawful import Enrico, Ursula
+from nucypher.crypto.powers import RitualisticPower
 from nucypher.policy.conditions.lingo import ConditionLingo, ConditionType
 from tests.constants import TESTERCHAIN_CHAIN_ID
 from tests.mock.coordinator import MockCoordinatorAgent
@@ -68,6 +69,9 @@ def cohort(ursulas, mock_coordinator_agent):
         mock_coordinator_agent._add_operator_to_staking_provider_mapping(
             {u.operator_address: u.checksum_address}
         )
+        mock_coordinator_agent.set_provider_public_key(
+            u.public_keys(RitualisticPower), u.transacting_power
+        )
         u.coordinator_agent = mock_coordinator_agent
         u.ritual_tracker.coordinator_agent = mock_coordinator_agent
 
diff --git a/tests/mock/coordinator.py b/tests/mock/coordinator.py
index 4ade040a7..756ce7d4b 100644
--- a/tests/mock/coordinator.py
+++ b/tests/mock/coordinator.py
@@ -170,9 +170,8 @@ class MockCoordinatorAgent(MockContractAgent):
         ritual.total_aggregations += 1
         return self.blockchain.FAKE_RECEIPT
 
-    @staticmethod
-    def is_provider_public_key_set(staking_provider: ChecksumAddress) -> bool:
-        return False
+    def is_provider_public_key_set(self, staking_provider: ChecksumAddress) -> bool:
+        return staking_provider in self._participant_keys_history
 
     def set_provider_public_key(
         self, public_key: FerveoPublicKey, transacting_power: TransactingPower
@@ -187,6 +186,7 @@ class MockCoordinatorAgent(MockContractAgent):
         participant_keys = self._participant_keys_history.get(provider_address)
         if not participant_keys:
             participant_keys = []
+            self._participant_keys_history[provider_address] = participant_keys
 
         participant_keys.append(
             self.ParticipantKey(
@@ -274,7 +274,7 @@ class MockCoordinatorAgent(MockContractAgent):
     def get_provider_public_key(
         self, provider: ChecksumAddress, ritual_id: int
     ) -> FerveoPublicKey:
-        participant_keys = self._participant_keys_history.get(provider)
+        participant_keys = self._participant_keys_history[provider]
         for participant_key in reversed(participant_keys):
             if participant_key.lastRitualId <= ritual_id:
                 g2Point = participant_key.publicKey
diff --git a/tests/unit/test_ritualist.py b/tests/unit/test_ritualist.py
index f4a754bab..3143d348d 100644
--- a/tests/unit/test_ritualist.py
+++ b/tests/unit/test_ritualist.py
@@ -2,16 +2,23 @@ import pytest
 
 from nucypher.blockchain.eth.agents import CoordinatorAgent
 from nucypher.blockchain.eth.signers.software import Web3Signer
-from nucypher.crypto.powers import TransactingPower
+from nucypher.crypto.powers import RitualisticPower, TransactingPower
 from tests.constants import MOCK_ETH_PROVIDER_URI
 from tests.mock.coordinator import MockCoordinatorAgent
 
 
 @pytest.fixture(scope="module")
-def agent(mock_contract_agency) -> MockCoordinatorAgent:
+def agent(mock_contract_agency, ursulas) -> MockCoordinatorAgent:
     coordinator_agent: CoordinatorAgent = mock_contract_agency.get_agent(
         CoordinatorAgent, registry=None, blockchain_endpoint=MOCK_ETH_PROVIDER_URI
     )
+
+    def mock_get_provider_public_key(provider, ritual_id):
+        for ursula in ursulas:
+            if ursula.checksum_address == provider:
+                return ursula.public_keys(RitualisticPower)
+
+    coordinator_agent.get_provider_public_key = mock_get_provider_public_key
     return coordinator_agent