From aedcab846ca3f2cac2e5215e8da48dfc9f5ba302 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Sat, 24 Jul 2021 19:15:51 +0000
Subject: [PATCH] RpcClient::send<T> now supports client-defined RPC methods
 via RpcRequest::Custom (#18882)

(cherry picked from commit f26451158540d9ef4103b3f2c905df8fddeaa408)

Co-authored-by: Michael Vines <mvines@gmail.com>
---
 client/src/mock_sender.rs | 33 ++++++++++++++++++---------------
 client/src/rpc_client.rs  | 15 +++++++++++++++
 client/src/rpc_request.rs |  4 ++++
 3 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/client/src/mock_sender.rs b/client/src/mock_sender.rs
index e40e456745546e..8af926bdfed268 100644
--- a/client/src/mock_sender.rs
+++ b/client/src/mock_sender.rs
@@ -49,23 +49,26 @@ impl RpcSender for MockSender {
         if self.url == "fails" {
             return Ok(Value::Null);
         }
-        let val = match request {
-            RpcRequest::GetAccountInfo => serde_json::to_value(Response {
+
+        let method = &request.build_request_json(42, params.clone())["method"];
+
+        let val = match method.as_str().unwrap() {
+            "getAccountInfo" => serde_json::to_value(Response {
                 context: RpcResponseContext { slot: 1 },
                 value: Value::Null,
             })?,
-            RpcRequest::GetBalance => serde_json::to_value(Response {
+            "getBalance" => serde_json::to_value(Response {
                 context: RpcResponseContext { slot: 1 },
                 value: Value::Number(Number::from(50)),
             })?,
-            RpcRequest::GetRecentBlockhash => serde_json::to_value(Response {
+            "getRecentBlockhash" => serde_json::to_value(Response {
                 context: RpcResponseContext { slot: 1 },
                 value: (
                     Value::String(PUBKEY.to_string()),
                     serde_json::to_value(FeeCalculator::default()).unwrap(),
                 ),
             })?,
-            RpcRequest::GetEpochInfo => serde_json::to_value(EpochInfo {
+            "getEpochInfo" => serde_json::to_value(EpochInfo {
                 epoch: 1,
                 slot_index: 2,
                 slots_in_epoch: 32,
@@ -73,7 +76,7 @@ impl RpcSender for MockSender {
                 block_height: 34,
                 transaction_count: Some(123),
             })?,
-            RpcRequest::GetFeeCalculatorForBlockhash => {
+            "getFeeCalculatorForBlockhash" => {
                 let value = if self.url == "blockhash_expired" {
                     Value::Null
                 } else {
@@ -84,11 +87,11 @@ impl RpcSender for MockSender {
                     value,
                 })?
             }
-            RpcRequest::GetFeeRateGovernor => serde_json::to_value(Response {
+            "getFeeRateGovernor" => serde_json::to_value(Response {
                 context: RpcResponseContext { slot: 1 },
                 value: serde_json::to_value(FeeRateGovernor::default()).unwrap(),
             })?,
-            RpcRequest::GetSignatureStatuses => {
+            "getSignatureStatuses" => {
                 let status: transaction::Result<()> = if self.url == "account_in_use" {
                     Err(TransactionError::AccountInUse)
                 } else if self.url == "instruction_error" {
@@ -122,11 +125,11 @@ impl RpcSender for MockSender {
                     value: statuses,
                 })?
             }
-            RpcRequest::GetTransactionCount => Value::Number(Number::from(1234)),
-            RpcRequest::GetSlot => Value::Number(Number::from(0)),
-            RpcRequest::GetMaxShredInsertSlot => Value::Number(Number::from(0)),
-            RpcRequest::RequestAirdrop => Value::String(Signature::new(&[8; 64]).to_string()),
-            RpcRequest::SendTransaction => {
+            "getTransactionCount" => json![1234],
+            "getSlot" => json![0],
+            "getMaxShredInsertSlot" => json![0],
+            "requestAirdrop" => Value::String(Signature::new(&[8; 64]).to_string()),
+            "sendTransaction" => {
                 let signature = if self.url == "malicious" {
                     Signature::new(&[8; 64]).to_string()
                 } else {
@@ -137,8 +140,8 @@ impl RpcSender for MockSender {
                 };
                 Value::String(signature)
             }
-            RpcRequest::GetMinimumBalanceForRentExemption => Value::Number(Number::from(20)),
-            RpcRequest::GetVersion => {
+            "getMinimumBalanceForRentExemption" => json![20],
+            "getVersion" => {
                 let version = Version::default();
                 json!(RpcVersionInfo {
                     solana_core: version.to_string(),
diff --git a/client/src/rpc_client.rs b/client/src/rpc_client.rs
index 642887a811acc6..97bc4427834e6f 100644
--- a/client/src/rpc_client.rs
+++ b/client/src/rpc_client.rs
@@ -2162,6 +2162,7 @@ mod tests {
         let signature = rpc_client.send_transaction(&tx);
         assert!(signature.is_err());
     }
+
     #[test]
     fn test_get_recent_blockhash() {
         let rpc_client = RpcClient::new_mock("succeeds".to_string());
@@ -2176,6 +2177,20 @@ mod tests {
         assert!(rpc_client.get_recent_blockhash().is_err());
     }
 
+    #[test]
+    fn test_custom_request() {
+        let rpc_client = RpcClient::new_mock("succeeds".to_string());
+
+        let slot = rpc_client.get_slot().unwrap();
+        assert_eq!(slot, 0);
+
+        let custom_slot = rpc_client
+            .send::<Slot>(RpcRequest::Custom { method: "getSlot" }, Value::Null)
+            .unwrap();
+
+        assert_eq!(slot, custom_slot);
+    }
+
     #[test]
     fn test_get_signature_status() {
         let signature = Signature::default();
diff --git a/client/src/rpc_request.rs b/client/src/rpc_request.rs
index 25d76ff28404be..52bae09c161024 100644
--- a/client/src/rpc_request.rs
+++ b/client/src/rpc_request.rs
@@ -86,6 +86,9 @@ pub enum RpcRequest {
     SendTransaction,
     SimulateTransaction,
     SignVote,
+    Custom {
+        method: &'static str,
+    },
 }
 
 #[allow(deprecated)]
@@ -154,6 +157,7 @@ impl fmt::Display for RpcRequest {
             RpcRequest::SendTransaction => "sendTransaction",
             RpcRequest::SimulateTransaction => "simulateTransaction",
             RpcRequest::SignVote => "signVote",
+            RpcRequest::Custom { method } => method,
         };
 
         write!(f, "{}", method)