From 01b6815dec85b67d25e1b6b1a436e83f03e8d631 Mon Sep 17 00:00:00 2001
From: brentstone <brent@heliax.dev>
Date: Tue, 3 Oct 2023 22:40:25 -0400
Subject: [PATCH] fix tests

---
 .../shell/vote_extensions/bridge_pool_vext.rs | 16 +-------
 proof_of_stake/src/pos_queries.rs             | 37 +++++++++----------
 2 files changed, 18 insertions(+), 35 deletions(-)

diff --git a/apps/src/lib/node/ledger/shell/vote_extensions/bridge_pool_vext.rs b/apps/src/lib/node/ledger/shell/vote_extensions/bridge_pool_vext.rs
index 4770fc67376..05f0263a10e 100644
--- a/apps/src/lib/node/ledger/shell/vote_extensions/bridge_pool_vext.rs
+++ b/apps/src/lib/node/ledger/shell/vote_extensions/bridge_pool_vext.rs
@@ -276,7 +276,6 @@ mod test_bp_vote_extensions {
     #[cfg(not(feature = "abcipp"))]
     use namada::ledger::eth_bridge::EthBridgeQueries;
     use namada::ledger::pos::PosQueries;
-    use namada::ledger::storage_api::StorageWrite;
     use namada::proof_of_stake::types::{
         Position as ValidatorPosition, WeightedValidator,
     };
@@ -321,25 +320,12 @@ mod test_bp_vote_extensions {
             )
             .expect("Test failed");
 
-        // register Bertha's protocol key
-        let pk_key = protocol_pk_key(&bertha_address());
-        shell
-            .wl_storage
-            .write_bytes(
-                &pk_key,
-                bertha_keypair()
-                    .ref_to()
-                    .try_to_vec()
-                    .expect("Test failed."),
-            )
-            .expect("Test failed.");
-
         // change pipeline length to 1
         let mut params = shell.wl_storage.pos_queries().get_pos_params();
         params.owned.pipeline_len = 1;
 
         let consensus_key = gen_keypair();
-        let protocol_key = gen_keypair();
+        let protocol_key = bertha_keypair();
         let hot_key = gen_secp256k1_keypair();
         let cold_key = gen_secp256k1_keypair();
 
diff --git a/proof_of_stake/src/pos_queries.rs b/proof_of_stake/src/pos_queries.rs
index 190548570ba..b694897aa86 100644
--- a/proof_of_stake/src/pos_queries.rs
+++ b/proof_of_stake/src/pos_queries.rs
@@ -1,7 +1,6 @@
 //! Storage API for querying data about Proof-of-stake related
 //! data. This includes validator and epoch related data.
 
-use borsh::{BorshDeserialize, BorshSerialize};
 use namada_core::ledger::parameters::storage::get_max_proposal_bytes_key;
 use namada_core::ledger::parameters::EpochDuration;
 use namada_core::ledger::storage::WlStorage;
@@ -172,17 +171,17 @@ where
         pk: &key::common::PublicKey,
         epoch: Option<Epoch>,
     ) -> Result<WeightedValidator> {
-        let pk_bytes = pk
-            .try_to_vec()
-            .expect("Serializing public key should not fail");
+        let params = crate::read_pos_params(self.wl_storage)
+            .expect("Failed to fetch Pos params");
         let epoch = epoch
             .unwrap_or_else(|| self.wl_storage.storage.get_current_epoch().0);
         self.get_consensus_validators(Some(epoch))
             .iter()
             .find(|validator| {
-                let pk_key = key::protocol_pk_key(&validator.address);
-                match self.wl_storage.storage.read(&pk_key) {
-                    Ok((Some(bytes), _)) => bytes == pk_bytes,
+                let protocol_keys =
+                    crate::validator_protocol_key_handle(&validator.address);
+                match protocol_keys.get(self.wl_storage, epoch, &params) {
+                    Ok(Some(key)) => key == *pk,
                     _ => false,
                 }
             })
@@ -195,26 +194,24 @@ where
         address: &Address,
         epoch: Option<Epoch>,
     ) -> Result<(token::Amount, key::common::PublicKey)> {
+        let params = crate::read_pos_params(self.wl_storage)
+            .expect("Failed to fetch Pos params");
         let epoch = epoch
             .unwrap_or_else(|| self.wl_storage.storage.get_current_epoch().0);
         self.get_consensus_validators(Some(epoch))
             .iter()
             .find(|validator| address == &validator.address)
             .map(|validator| {
-                let protocol_pk_key = key::protocol_pk_key(&validator.address);
-                // TODO: rewrite this, to use `StorageRead::read`
-                let bytes = self
-                    .wl_storage
-                    .storage
-                    .read(&protocol_pk_key)
-                    .expect("Validator should have public protocol key")
-                    .0
-                    .expect("Validator should have public protocol key");
-                let protocol_pk: key::common::PublicKey =
-                    BorshDeserialize::deserialize(&mut bytes.as_ref()).expect(
-                        "Protocol public key in storage should be \
-                         deserializable",
+                let protocol_keys =
+                    crate::validator_protocol_key_handle(&validator.address);
+                let protocol_pk = protocol_keys
+                    .get(self.wl_storage, epoch, &params)
+                    .unwrap()
+                    .expect(
+                        "Protocol public key should be set in storage after \
+                         genesis.",
                     );
+
                 (validator.bonded_stake, protocol_pk)
             })
             .ok_or_else(|| Error::NotValidatorAddress(address.clone(), epoch))