From 7ddecea7b5f081b530ea86d3798c7c47c9f8602c Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Fri, 23 Jul 2021 14:39:38 -0700
Subject: [PATCH] RpcClient::send<T> now supports client-defined RPC methods
 via RpcRequest::Custom

---
 client/src/mock_sender.rs | 35 +++++++++++++++++++----------------
 client/src/rpc_client.rs  | 15 +++++++++++++++
 client/src/rpc_request.rs |  4 ++++
 3 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/client/src/mock_sender.rs b/client/src/mock_sender.rs
index 95cb28f218f310..c13caaaf48ae82 100644
--- a/client/src/mock_sender.rs
+++ b/client/src/mock_sender.rs
@@ -78,23 +78,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,
@@ -102,7 +105,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 {
@@ -113,11 +116,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" {
@@ -151,11 +154,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 {
@@ -166,7 +169,7 @@ impl RpcSender for MockSender {
                 };
                 Value::String(signature)
             }
-            RpcRequest::SimulateTransaction => serde_json::to_value(Response {
+            "simulateTransaction" => serde_json::to_value(Response {
                 context: RpcResponseContext { slot: 1 },
                 value: RpcSimulateTransactionResult {
                     err: None,
@@ -175,8 +178,8 @@ impl RpcSender for MockSender {
                     units_consumed: None,
                 },
             })?,
-            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 ec86c2e03c80f8..7968e69776adc2 100644
--- a/client/src/rpc_client.rs
+++ b/client/src/rpc_client.rs
@@ -2575,6 +2575,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());
@@ -2589,6 +2590,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)