1use std::{
7 collections::{HashMap, VecDeque},
8 mem,
9 sync::Mutex,
10};
11
12use linera_base::{
13 abi::ServiceAbi,
14 data_types::{Amount, BlockHeight, Timestamp},
15 hex, http,
16 identifiers::{AccountOwner, ApplicationId, ChainId, DataBlobHash},
17};
18use serde::{de::DeserializeOwned, Serialize};
19
20use crate::{KeyValueStore, Service, ViewStorageContext};
21
22pub struct MockServiceRuntime<Application>
24where
25 Application: Service,
26{
27 application_parameters: Mutex<Option<Application::Parameters>>,
28 application_id: Mutex<Option<ApplicationId<Application::Abi>>>,
29 application_creator_chain_id: Mutex<Option<ChainId>>,
30 chain_id: Mutex<Option<ChainId>>,
31 next_block_height: Mutex<Option<BlockHeight>>,
32 timestamp: Mutex<Option<Timestamp>>,
33 chain_balance: Mutex<Option<Amount>>,
34 owner_balances: Mutex<Option<HashMap<AccountOwner, Amount>>>,
35 query_application_handler: Mutex<Option<QueryApplicationHandler>>,
36 expected_http_requests: Mutex<VecDeque<(http::Request, http::Response)>>,
37 blobs: Mutex<Option<HashMap<DataBlobHash, Vec<u8>>>>,
38 scheduled_operations: Mutex<Vec<Vec<u8>>>,
39 key_value_store: KeyValueStore,
40}
41
42impl<Application> Default for MockServiceRuntime<Application>
43where
44 Application: Service,
45{
46 fn default() -> Self {
47 MockServiceRuntime::new()
48 }
49}
50
51impl<Application> MockServiceRuntime<Application>
52where
53 Application: Service,
54{
55 pub fn new() -> Self {
57 MockServiceRuntime {
58 application_parameters: Mutex::new(None),
59 application_id: Mutex::new(None),
60 application_creator_chain_id: Mutex::new(None),
61 chain_id: Mutex::new(None),
62 next_block_height: Mutex::new(None),
63 timestamp: Mutex::new(None),
64 chain_balance: Mutex::new(None),
65 owner_balances: Mutex::new(None),
66 query_application_handler: Mutex::new(None),
67 expected_http_requests: Mutex::new(VecDeque::new()),
68 blobs: Mutex::new(None),
69 scheduled_operations: Mutex::new(vec![]),
70 key_value_store: KeyValueStore::mock(),
71 }
72 }
73
74 pub fn key_value_store(&self) -> KeyValueStore {
76 self.key_value_store.clone()
77 }
78
79 pub fn root_view_storage_context(&self) -> ViewStorageContext {
81 ViewStorageContext::new_unsafe(self.key_value_store(), Vec::new(), ())
82 }
83
84 pub fn with_application_parameters(
86 self,
87 application_parameters: Application::Parameters,
88 ) -> Self {
89 *self.application_parameters.lock().unwrap() = Some(application_parameters);
90 self
91 }
92
93 pub fn set_application_parameters(
95 &self,
96 application_parameters: Application::Parameters,
97 ) -> &Self {
98 *self.application_parameters.lock().unwrap() = Some(application_parameters);
99 self
100 }
101
102 pub fn application_parameters(&self) -> Application::Parameters {
104 Self::fetch_mocked_value(
105 &self.application_parameters,
106 "Application parameters have not been mocked, \
107 please call `MockServiceRuntime::set_application_parameters` first",
108 )
109 }
110
111 pub fn with_application_id(self, application_id: ApplicationId<Application::Abi>) -> Self {
113 *self.application_id.lock().unwrap() = Some(application_id);
114 self
115 }
116
117 pub fn set_application_id(&self, application_id: ApplicationId<Application::Abi>) -> &Self {
119 *self.application_id.lock().unwrap() = Some(application_id);
120 self
121 }
122
123 pub fn application_id(&self) -> ApplicationId<Application::Abi> {
125 Self::fetch_mocked_value(
126 &self.application_id,
127 "Application ID has not been mocked, \
128 please call `MockServiceRuntime::set_application_id` first",
129 )
130 }
131
132 pub fn with_application_creator_chain_id(self, application_creator_chain_id: ChainId) -> Self {
134 *self.application_creator_chain_id.lock().unwrap() = Some(application_creator_chain_id);
135 self
136 }
137
138 pub fn set_application_creator_chain_id(&self, application_creator_chain_id: ChainId) -> &Self {
140 *self.application_creator_chain_id.lock().unwrap() = Some(application_creator_chain_id);
141 self
142 }
143
144 pub fn application_creator_chain_id(&self) -> ChainId {
146 Self::fetch_mocked_value(
147 &self.application_creator_chain_id,
148 "Application creator chain ID has not been mocked, \
149 please call `MockServiceRuntime::set_application_creator_chain_id` first",
150 )
151 }
152
153 pub fn with_chain_id(self, chain_id: ChainId) -> Self {
155 *self.chain_id.lock().unwrap() = Some(chain_id);
156 self
157 }
158
159 pub fn set_chain_id(&self, chain_id: ChainId) -> &Self {
161 *self.chain_id.lock().unwrap() = Some(chain_id);
162 self
163 }
164
165 pub fn chain_id(&self) -> ChainId {
167 Self::fetch_mocked_value(
168 &self.chain_id,
169 "Chain ID has not been mocked, \
170 please call `MockServiceRuntime::set_chain_id` first",
171 )
172 }
173
174 pub fn with_next_block_height(self, next_block_height: BlockHeight) -> Self {
176 *self.next_block_height.lock().unwrap() = Some(next_block_height);
177 self
178 }
179
180 pub fn set_next_block_height(&self, next_block_height: BlockHeight) -> &Self {
182 *self.next_block_height.lock().unwrap() = Some(next_block_height);
183 self
184 }
185
186 pub fn next_block_height(&self) -> BlockHeight {
188 Self::fetch_mocked_value(
189 &self.next_block_height,
190 "Next block height has not been mocked, \
191 please call `MockServiceRuntime::set_next_block_height` first",
192 )
193 }
194
195 pub fn with_system_time(self, timestamp: Timestamp) -> Self {
197 *self.timestamp.lock().unwrap() = Some(timestamp);
198 self
199 }
200
201 pub fn set_system_time(&self, timestamp: Timestamp) -> &Self {
203 *self.timestamp.lock().unwrap() = Some(timestamp);
204 self
205 }
206
207 pub fn system_time(&self) -> Timestamp {
209 Self::fetch_mocked_value(
210 &self.timestamp,
211 "System time has not been mocked, \
212 please call `MockServiceRuntime::set_system_time` first",
213 )
214 }
215
216 pub fn with_chain_balance(self, chain_balance: Amount) -> Self {
218 *self.chain_balance.lock().unwrap() = Some(chain_balance);
219 self
220 }
221
222 pub fn set_chain_balance(&self, chain_balance: Amount) -> &Self {
224 *self.chain_balance.lock().unwrap() = Some(chain_balance);
225 self
226 }
227
228 pub fn chain_balance(&self) -> Amount {
230 Self::fetch_mocked_value(
231 &self.chain_balance,
232 "Chain balance has not been mocked, \
233 please call `MockServiceRuntime::set_chain_balance` first",
234 )
235 }
236
237 pub fn with_owner_balances(
239 self,
240 owner_balances: impl IntoIterator<Item = (AccountOwner, Amount)>,
241 ) -> Self {
242 *self.owner_balances.lock().unwrap() = Some(owner_balances.into_iter().collect());
243 self
244 }
245
246 pub fn set_owner_balances(
248 &self,
249 owner_balances: impl IntoIterator<Item = (AccountOwner, Amount)>,
250 ) -> &Self {
251 *self.owner_balances.lock().unwrap() = Some(owner_balances.into_iter().collect());
252 self
253 }
254
255 pub fn with_owner_balance(self, owner: AccountOwner, balance: Amount) -> Self {
257 self.set_owner_balance(owner, balance);
258 self
259 }
260
261 pub fn set_owner_balance(&self, owner: AccountOwner, balance: Amount) -> &Self {
263 self.owner_balances
264 .lock()
265 .unwrap()
266 .get_or_insert_with(HashMap::new)
267 .insert(owner, balance);
268 self
269 }
270
271 pub fn owner_balance(&self, owner: AccountOwner) -> Amount {
273 self.owner_balances
274 .lock()
275 .unwrap()
276 .as_mut()
277 .and_then(|owner_balances| owner_balances.get(&owner).copied())
278 .unwrap_or_else(|| {
279 panic!(
280 "Balance for owner {owner} was not mocked, \
281 please include a balance for them with a call to \
282 `MockServiceRuntime::set_owner_balance`"
283 )
284 })
285 }
286
287 pub fn owner_balances(&self) -> Vec<(AccountOwner, Amount)> {
289 self.owner_balances
290 .lock()
291 .unwrap()
292 .as_ref()
293 .expect(
294 "Owner balances have not been mocked, \
295 please call `MockServiceRuntime::set_owner_balances` first",
296 )
297 .iter()
298 .map(|(owner, amount)| (*owner, *amount))
299 .collect()
300 }
301
302 pub fn balance_owners(&self) -> Vec<AccountOwner> {
304 self.owner_balances
305 .lock()
306 .unwrap()
307 .as_ref()
308 .expect(
309 "Owner balances have not been mocked, \
310 please call `MockServiceRuntime::set_owner_balances` first",
311 )
312 .keys()
313 .copied()
314 .collect()
315 }
316
317 pub fn schedule_raw_operation(&self, operation: Vec<u8>) {
321 self.scheduled_operations.lock().unwrap().push(operation);
322 }
323
324 pub fn schedule_operation(&self, operation: &impl Serialize) {
328 let bytes = bcs::to_bytes(operation).expect("Failed to serialize application operation");
329
330 self.schedule_raw_operation(bytes);
331 }
332
333 pub fn raw_scheduled_operations(&self) -> Vec<Vec<u8>> {
339 mem::take(&mut self.scheduled_operations.lock().unwrap())
340 }
341
342 pub fn scheduled_operations<Operation>(&self) -> Vec<Operation>
350 where
351 Operation: DeserializeOwned,
352 {
353 self.raw_scheduled_operations()
354 .into_iter()
355 .enumerate()
356 .map(|(index, bytes)| {
357 bcs::from_bytes(&bytes).unwrap_or_else(|error| {
358 panic!(
359 "Failed to deserialize scheduled operation #{index} (0x{}): {error}",
360 hex::encode(bytes)
361 )
362 })
363 })
364 .collect()
365 }
366
367 pub fn with_query_application_handler(
369 self,
370 handler: impl FnMut(ApplicationId, Vec<u8>) -> Vec<u8> + Send + 'static,
371 ) -> Self {
372 *self.query_application_handler.lock().unwrap() = Some(Box::new(handler));
373 self
374 }
375
376 pub fn set_query_application_handler(
378 &self,
379 handler: impl FnMut(ApplicationId, Vec<u8>) -> Vec<u8> + Send + 'static,
380 ) -> &Self {
381 *self.query_application_handler.lock().unwrap() = Some(Box::new(handler));
382 self
383 }
384
385 pub fn query_application<A: ServiceAbi>(
387 &self,
388 application: ApplicationId<A>,
389 query: &A::Query,
390 ) -> A::QueryResponse {
391 let query_bytes =
392 serde_json::to_vec(&query).expect("Failed to serialize query to another application");
393
394 let mut handler_guard = self.query_application_handler.lock().unwrap();
395 let handler = handler_guard.as_mut().expect(
396 "Handler for `query_application` has not been mocked, \
397 please call `MockServiceRuntime::set_query_application_handler` first",
398 );
399
400 let response_bytes = handler(application.forget_abi(), query_bytes);
401
402 serde_json::from_slice(&response_bytes)
403 .expect("Failed to deserialize query response from application")
404 }
405
406 pub fn add_expected_http_request(&mut self, request: http::Request, response: http::Response) {
408 self.expected_http_requests
409 .lock()
410 .unwrap()
411 .push_back((request, response));
412 }
413
414 pub fn http_request(&self, request: http::Request) -> http::Response {
422 let maybe_request = self.expected_http_requests.lock().unwrap().pop_front();
423 let (expected_request, response) = maybe_request.expect("Unexpected HTTP request");
424 assert_eq!(request, expected_request);
425 response
426 }
427
428 pub fn with_blobs(self, blobs: impl IntoIterator<Item = (DataBlobHash, Vec<u8>)>) -> Self {
430 *self.blobs.lock().unwrap() = Some(blobs.into_iter().collect());
431 self
432 }
433
434 pub fn set_blobs(&self, blobs: impl IntoIterator<Item = (DataBlobHash, Vec<u8>)>) -> &Self {
436 *self.blobs.lock().unwrap() = Some(blobs.into_iter().collect());
437 self
438 }
439
440 pub fn with_blob(self, hash: impl Into<DataBlobHash>, blob: Vec<u8>) -> Self {
442 self.set_blob(hash, blob);
443 self
444 }
445
446 pub fn set_blob(&self, hash: impl Into<DataBlobHash>, blob: Vec<u8>) -> &Self {
448 self.blobs
449 .lock()
450 .unwrap()
451 .get_or_insert_with(HashMap::new)
452 .insert(hash.into(), blob);
453 self
454 }
455
456 pub fn read_data_blob(&self, hash: DataBlobHash) -> Vec<u8> {
458 self.blobs
459 .lock()
460 .unwrap()
461 .as_ref()
462 .and_then(|blobs| blobs.get(&hash).cloned())
463 .unwrap_or_else(|| {
464 panic!(
465 "Blob for hash {hash:?} has not been mocked, \
466 please call `MockServiceRuntime::set_blob` first"
467 )
468 })
469 }
470
471 pub fn assert_blob_exists(&self, hash: DataBlobHash) {
473 assert!(
474 self.blobs
475 .lock()
476 .unwrap()
477 .as_ref()
478 .is_some_and(|blobs| blobs.contains_key(&hash)),
479 "Blob for hash {hash:?} has not been mocked, \
480 please call `MockServiceRuntime::set_blob` first"
481 );
482 }
483
484 fn fetch_mocked_value<T>(slot: &Mutex<Option<T>>, message: &str) -> T
486 where
487 T: Clone,
488 {
489 slot.lock().unwrap().clone().expect(message)
490 }
491}
492
493pub type QueryApplicationHandler = Box<dyn FnMut(ApplicationId, Vec<u8>) -> Vec<u8> + Send>;