neo3/neo_clients/
mock_client.rs

1use crate::{
2	crypto::{KeyPair, Secp256r1PrivateKey},
3	neo_clients::{HttpProvider, RpcClient},
4	neo_config::TestConstants,
5	neo_protocol::{Account, AccountTrait, ApplicationLog, NeoVersion, RawTransaction},
6};
7use lazy_static::lazy_static;
8use neo3::prelude::*;
9use primitive_types::{H160, H256};
10use regex::Regex;
11use serde_json::{json, Value};
12use std::{fs, path::PathBuf, str::FromStr, sync::Arc};
13use tokio::sync::Mutex;
14use url::Url;
15use wiremock::{
16	matchers::{body_json, body_partial_json, method, path},
17	Match, Mock, MockServer, ResponseTemplate,
18};
19
20lazy_static! {
21	pub static ref ACCOUNT1: Account = Account::from_key_pair(
22		KeyPair::from_secret_key(
23			&Secp256r1PrivateKey::from_bytes(
24				&hex::decode("e6e919577dd7b8e97805151c05ae07ff4f752654d6d8797597aca989c02c4cb3")
25					.unwrap()
26			)
27			.unwrap()
28		),
29		None,
30		None
31	)
32	.expect("Failed to create ACCOUNT1");
33	pub static ref ACCOUNT2: Account = Account::from_key_pair(
34		KeyPair::from_secret_key(
35			&Secp256r1PrivateKey::from_bytes(
36				&hex::decode("b4b2b579cac270125259f08a5f414e9235817e7637b9a66cfeb3b77d90c8e7f9")
37					.unwrap()
38			)
39			.unwrap()
40		),
41		None,
42		None
43	)
44	.expect("Failed to create ACCOUNT2");
45}
46
47pub struct MockClient {
48	server: MockServer,
49	mocks: Vec<Mock>,
50}
51
52impl MockClient {
53	pub async fn new() -> Self {
54		let server = MockServer::start().await;
55		Self { server, mocks: Vec::new() }
56	}
57
58	pub async fn mock_response(
59		&mut self,
60		method_name: &str,
61		params: serde_json::Value,
62		result: serde_json::Value,
63	) {
64		let mock = Mock::given(method("POST"))
65			.and(path("/"))
66			.and(body_json(json!({
67				"jsonrpc": "2.0",
68				"method": method_name,
69				"params": params,
70				"id": 1
71			})))
72			.respond_with(ResponseTemplate::new(200).set_body_json(json!({
73				"jsonrpc": "2.0",
74				"id": 1,
75				"result": result
76			})));
77		self.mocks.push(mock);
78	}
79
80	pub async fn mock_response_error(&mut self, error: serde_json::Value) {
81		let mock = Mock::given(method("POST")).and(path("/")).respond_with(
82			ResponseTemplate::new(200).set_body_json(json!({
83				"jsonrpc": "2.0",
84				"id": 1,
85				"error": error
86			})),
87		);
88		self.mocks.push(mock);
89	}
90
91	pub async fn mock_response_ignore_param(
92		&mut self,
93		method_name: &str,
94		result: serde_json::Value,
95	) -> &mut Self {
96		let mock = Mock::given(method("POST"))
97			.and(path("/"))
98			.and(body_partial_json(json!({
99				"jsonrpc": "2.0",
100				"method": method_name,
101			})))
102			.respond_with(ResponseTemplate::new(200).set_body_json(json!({
103				"jsonrpc": "2.0",
104				"id": 1,
105				"result": result
106			})));
107		self.mocks.push(mock);
108		self
109	}
110
111	pub async fn mock_response_with_file(
112		&mut self,
113		method_name: &str,
114		response_file: &str,
115		params: serde_json::Value,
116	) -> &mut Self {
117		// Construct the path to the response file relative to the project root
118		let mut response_file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
119		response_file_path.push("test_resources");
120		response_file_path.push("responses");
121		response_file_path.push(response_file);
122
123		// Load the response body from the specified file
124		let response_body = tokio::fs::read_to_string(response_file_path)
125			.await
126			.expect("Failed to read response file");
127
128		let mock = Mock::given(method("POST"))
129			.and(path("/"))
130			.and(body_partial_json(json!({
131				"jsonrpc": "2.0",
132				"method": method_name,
133				"params": params,
134			})))
135			.respond_with(ResponseTemplate::new(200).set_body_string(response_body));
136		self.mocks.push(mock);
137		self
138	}
139
140	pub async fn mock_response_with_file_ignore_param(
141		&mut self,
142		method_name: &str,
143		response_file: &str,
144	) -> &mut Self {
145		// Construct the path to the response file relative to the project root
146		let mut response_file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
147		response_file_path.push("test_resources");
148		response_file_path.push("responses");
149		response_file_path.push(response_file);
150
151		// Load the response body from the specified file
152		let response_body = tokio::fs::read_to_string(response_file_path)
153			.await
154			.expect("Failed to read response file");
155		// // Load the response body from the specified file
156		// let response_body = tokio::fs::read_to_string(format!("/responses/{}", response_file))
157		// .await
158		// .expect("Failed to read response file");
159
160		let mock = Mock::given(method("POST"))
161			.and(path("/"))
162			.and(body_partial_json(json!({
163				"jsonrpc": "2.0",
164				"method": method_name,
165			})))
166			.respond_with(ResponseTemplate::new(200).set_body_string(response_body));
167		self.mocks.push(mock);
168		self
169	}
170
171	pub async fn mock_response_for_balance_of(
172		&mut self,
173		contract_hash: &str,
174		account_script_hash: &str,
175		response_file: &str,
176	) -> &mut Self {
177		// Construct the path to the response file relative to the project root
178		let mut response_file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
179		response_file_path.push("test_resources");
180		response_file_path.push("responses");
181		response_file_path.push(response_file);
182
183		// Load the response body from the specified file
184		let response_body = tokio::fs::read_to_string(response_file_path)
185			.await
186			.expect("Failed to read response file");
187
188		let mock = Mock::given(method("POST"))
189			.and(path("/"))
190			.and(body_partial_json(json!({
191				"jsonrpc": "2.0",
192				"method": "invokefunction",
193				"params": [
194					contract_hash,
195					"balanceOf",
196					[
197						{
198							"type": "Hash160",
199							"value": account_script_hash,
200						}
201					]
202				],
203			})))
204			.respond_with(ResponseTemplate::new(200).set_body_string(response_body));
205
206		self.mocks.push(mock);
207		self
208	}
209
210	pub async fn mock_default_responses(&mut self) -> &mut Self {
211		self.mock_response_with_file_ignore_param(
212			"invokescript",
213			"invokescript_necessary_mock.json",
214		)
215		.await;
216		self.mock_response_with_file(
217			"invokefunction",
218			"invokefunction_transfer_neo.json",
219			json!([
220				TestConstants::NEO_TOKEN_HASH,
221				"transfer",
222				vec![
223					ContractParameter::from(ACCOUNT1.address_or_scripthash().script_hash()),
224					ContractParameter::from(
225						H160::from_str("969a77db482f74ce27105f760efa139223431394").unwrap(),
226					),
227					ContractParameter::from(5),
228					ContractParameter::any(),
229				],
230			]),
231		)
232		.await;
233		self.mock_response_with_file_ignore_param("getblockcount", "getblockcount_1000.json")
234			.await;
235		self.mock_response_with_file_ignore_param(
236			"calculatenetworkfee",
237			"calculatenetworkfee.json",
238		)
239		.await;
240		self
241	}
242
243	pub async fn mock_invoke_script(&mut self, result: InvocationResult) -> &mut Self {
244		self.mock_response_ignore_param("invokescript", json!(Ok::<InvocationResult, ()>(result)))
245			.await;
246		self
247	}
248
249	// pub async fn mock_get_block_count(&mut self, result: i32) -> &mut Self {
250	// 	self.mock_response_ignore_param("getblockcount", json!(Ok::<i32, ()>(result)))
251	// 		.await;
252	// 	self
253	// }
254
255	pub async fn mock_get_block_count(&mut self, block_count: u32) -> &mut Self {
256		// Construct the path to the response file relative to the project root
257		let mut response_file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
258		response_file_path.push("test_resources");
259		response_file_path.push("responses");
260		response_file_path.push(format!("getblockcount_{}.json", block_count));
261
262		// Load the response body from the specified file
263		let response_body = tokio::fs::read_to_string(response_file_path)
264			.await
265			.expect("Failed to read response file");
266		// // Load the response body from the specified file
267		// let response_body = tokio::fs::read_to_string(format!("/responses/{}", response_file))
268		// .await
269		// .expect("Failed to read response file");
270
271		let mock = Mock::given(method("POST"))
272			.and(path("/"))
273			.and(body_partial_json(json!({
274				"jsonrpc": "2.0",
275				"method": "getblockcount",
276			})))
277			.respond_with(ResponseTemplate::new(200).set_body_string(response_body));
278		self.mocks.push(mock);
279		self
280	}
281
282	// pub async fn mock_calculate_network_fee(&mut self, result: i32) -> &mut Self {
283	// 	self.mock_response_ignore_param("calculatenetworkfee", json!(Ok::<i32, ()>(result)))
284	// 		.await;
285	// 	self
286	// }
287
288	pub async fn mock_send_raw_transaction(&mut self, result: RawTransaction) -> &mut Self {
289		self.mock_response_with_file_ignore_param("sendrawtransaction", "sendrawtransaction.json")
290			.await;
291		self
292	}
293
294	pub async fn mock_get_version(&mut self, result: NeoVersion) -> &mut Self {
295		self.mock_response_ignore_param("getversion", json!(Ok::<NeoVersion, ()>(result)))
296			.await;
297		self
298	}
299
300	pub async fn mock_invoke_function(&mut self, result: InvocationResult) -> &mut Self {
301		self.mock_response_ignore_param(
302			"invokefunction",
303			json!(Ok::<InvocationResult, ()>(result)),
304		)
305		.await;
306		self
307	}
308
309	pub async fn mock_get_application_log(&mut self, result: Option<ApplicationLog>) -> &mut Self {
310		self.mock_response_ignore_param("getapplicationlog", json!(result)).await;
311		self
312	}
313
314	pub async fn mount_mocks(&mut self) -> &mut Self {
315		for mock in self.mocks.drain(..) {
316			mock.mount(&self.server).await;
317		}
318		self
319	}
320
321	pub fn url(&self) -> Url {
322		Url::parse(&self.server.uri()).expect("Invalid mock server URL")
323	}
324
325	pub fn into_client(&self) -> RpcClient<HttpProvider> {
326		let http_provider = HttpProvider::new(self.url()).expect("Failed to create HTTP provider");
327		RpcClient::new(http_provider)
328	}
329
330	pub fn server(&self) -> &MockServer {
331		&self.server
332	}
333}