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 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 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 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 let response_body = tokio::fs::read_to_string(response_file_path)
153 .await
154 .expect("Failed to read response file");
155 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 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 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, block_count: u32) -> &mut Self {
256 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 let response_body = tokio::fs::read_to_string(response_file_path)
264 .await
265 .expect("Failed to read response file");
266 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_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}