[test_loop] ClientQueries trait to use Borrow<Client> instead of AsRe…

…f<Client> + AsRef<AccountId>
near · Jun 8, 2024 · 8c33367 · 8c33367
1 parent dc03a34
commit 8c33367
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 21 deletions.
diff --git a/chain/client/src/test_utils/test_loop.rs b/chain/client/src/test_utils/test_loop.rs
@@ -3,6 +3,8 @@ pub mod partial_witness_actor;
 pub mod sync_actor;
 pub mod sync_jobs_actor;
 
+use std::borrow::Borrow;
+
 use crate::client_actor::{ClientActorInner, ClientSenderForPartialWitnessMessage};
 use near_async::messaging::{CanSend, Handler, SendAsync};
 use near_async::test_loop::delay_sender::DelaySender;
@@ -220,22 +222,22 @@ pub trait ClientQueries {
     fn tracked_shards_for_each_client(&self) -> Vec<Vec<ShardId>>;
 }
 
-impl<Data: AsRef<Client> + AsRef<AccountId>> ClientQueries for Vec<Data> {
+impl<Data> ClientQueries for Vec<Data>
+where
+    Data: Borrow<Client>,
+{
     fn client_index_tracking_account(&self, account_id: &AccountId) -> usize {
-        let client: &Client = self[0].as_ref();
+        let client: &Client = self[0].borrow();
         let head = client.chain.head().unwrap();
         let shard_id =
             client.epoch_manager.account_id_to_shard_id(&account_id, &head.epoch_id).unwrap();
 
         for i in 0..self.len() {
-            let client: &Client = self[i].as_ref();
+            let client: &Client = self[i].borrow();
+            let account_id = client.validator_signer.as_ref().unwrap().validator_id();
             let tracks_shard = client
                 .epoch_manager
-                .cares_about_shard_from_prev_block(
-                    &head.prev_block_hash,
-                    &self[i].as_ref(),
-                    shard_id,
-                )
+                .cares_about_shard_from_prev_block(&head.prev_block_hash, account_id, shard_id)
                 .unwrap();
             if tracks_shard {
                 return i;
@@ -246,7 +248,7 @@ impl<Data: AsRef<Client> + AsRef<AccountId>> ClientQueries for Vec<Data> {
 
     fn runtime_query(&self, account_id: &AccountId, query: QueryRequest) -> QueryResponse {
         let client_index = self.client_index_tracking_account(account_id);
-        let client: &Client = self[client_index].as_ref();
+        let client: &Client = self[client_index].borrow();
         let head = client.chain.head().unwrap();
         let last_block = client.chain.get_block(&head.last_block_hash).unwrap();
         let shard_id =
@@ -299,27 +301,24 @@ impl<Data: AsRef<Client> + AsRef<AccountId>> ClientQueries for Vec<Data> {
 
     fn tx_outcome(&self, tx_hash: CryptoHash) -> FinalExecutionOutcomeView {
         // TODO: this does not work yet with single-shard tracking.
-        let client: &Client = self[0].as_ref();
+        let client: &Client = self[0].borrow();
         client.chain.get_final_transaction_result(&tx_hash).unwrap()
     }
 
     fn tracked_shards_for_each_client(&self) -> Vec<Vec<ShardId>> {
-        let client: &Client = self[0].as_ref();
+        let client: &Client = self[0].borrow();
         let head = client.chain.head().unwrap();
         let all_shard_ids = client.epoch_manager.shard_ids(&head.epoch_id).unwrap();
 
         let mut ret = Vec::new();
         for i in 0..self.len() {
-            let client: &Client = self[i].as_ref();
+            let client: &Client = self[i].borrow();
+            let account_id = client.validator_signer.as_ref().unwrap().validator_id();
             let mut tracked_shards = Vec::new();
             for shard_id in &all_shard_ids {
                 let tracks_shard = client
                     .epoch_manager
-                    .cares_about_shard_from_prev_block(
-                        &head.prev_block_hash,
-                        &self[i].as_ref(),
-                        *shard_id,
-                    )
+                    .cares_about_shard_from_prev_block(&head.prev_block_hash, account_id, *shard_id)
                     .unwrap();
                 if tracks_shard {
                     tracked_shards.push(*shard_id);

diff --git a/integration-tests/src/tests/client/features/multinode_stateless_validators.rs b/integration-tests/src/tests/client/features/multinode_stateless_validators.rs
@@ -87,6 +87,7 @@ use near_vm_runner::ContractRuntimeCache;
 use near_vm_runner::FilesystemContractRuntimeCache;
 use nearcore::state_sync::StateSyncDumper;
 use nearcore::NightshadeRuntime;
+use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex, RwLock};
 
@@ -109,8 +110,8 @@ impl AsMut<TestData> for TestData {
     }
 }
 
-impl AsRef<Client> for TestData {
-    fn as_ref(&self) -> &Client {
+impl Borrow<Client> for TestData {
+    fn borrow(&self) -> &Client {
         &self.client.client
     }
 }

diff --git a/integration-tests/src/tests/client/features/multinode_test_loop_example.rs b/integration-tests/src/tests/client/features/multinode_test_loop_example.rs
@@ -84,6 +84,7 @@ use near_vm_runner::ContractRuntimeCache;
 use near_vm_runner::FilesystemContractRuntimeCache;
 use nearcore::state_sync::StateSyncDumper;
 use nearcore::NightshadeRuntime;
+use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex, RwLock};
 
@@ -106,8 +107,8 @@ impl AsMut<TestData> for TestData {
     }
 }
 
-impl AsRef<Client> for TestData {
-    fn as_ref(&self) -> &Client {
+impl Borrow<Client> for TestData {
+    fn borrow(&self) -> &Client {
         &self.client.client
     }
 }