1use std::sync::LazyLock;
19
20use ahash::AHashMap;
21use jiff::civil::{Time, Weekday};
22use nautilus_core::python::{
23 clone_py_object, to_pynotimplemented_err, to_pytype_err, to_pyvalue_err,
24};
25use nautilus_model::{
26 data::Data,
27 identifiers::{InstrumentId, Venue},
28 instruments::{Instrument, InstrumentAny},
29 orderbook::OrderBook,
30 position::Position,
31 python::{data::data_to_pyobject, instruments::instrument_any_to_pyobject},
32 types::{Currency, Money},
33};
34use parking_lot::Mutex;
35use pyo3::{
36 PyClass,
37 prelude::*,
38 types::{PyDict, PyTuple},
39};
40use rust_decimal::Decimal;
41
42use crate::modules::{
43 AccountAdjustmentOutcome, CfdSwapModule, CfdSwapRate, ExchangeContext,
44 FXRolloverInterestModule, SimulationModule, SimulationModuleAny, SimulationModuleHandle,
45 SimulationModuleResult, fx_rollover::InterestRateRecord,
46};
47
48pub type SimulationModuleExtractor =
50 for<'py> fn(Python<'py>, &Bound<'py, PyAny>) -> PyResult<SimulationModuleHandle>;
51
52static SIMULATION_MODULE_EXTRACTORS: LazyLock<Mutex<AHashMap<usize, SimulationModuleExtractor>>> =
53 LazyLock::new(|| Mutex::new(AHashMap::new()));
54
55pub fn register_simulation_module_extractor<T: PyClass>(
63 py: Python<'_>,
64 extractor: SimulationModuleExtractor,
65) -> anyhow::Result<()> {
66 let type_object = py.get_type::<T>();
67 let type_id = type_object.as_ptr() as usize;
68 let type_name = type_object.name()?;
69 let mut extractors = SIMULATION_MODULE_EXTRACTORS.lock();
70 if let Some(registered) = extractors.get(&type_id) {
71 if std::ptr::fn_addr_eq(*registered, extractor) {
72 return Ok(());
73 }
74 anyhow::bail!(
75 "A different simulation module extractor is already registered for '{type_name}'"
76 );
77 }
78 extractors.insert(type_id, extractor);
79 Ok(())
80}
81
82#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.backtest")]
83#[pyclass(
84 module = "nautilus_trader.backtest",
85 name = "SimulationModule",
86 subclass
87)]
88#[derive(Debug)]
89pub struct PySimulationModule;
90
91#[pyo3_stub_gen::derive::gen_stub_pymethods]
92#[pymethods]
93#[allow(
94 clippy::unused_self,
95 reason = "PyO3 exposes these hooks as overridable instance methods"
96)]
97impl PySimulationModule {
98 #[new]
99 #[gen_stub(override_return_type(type_repr = "typing.Self", imports = ("typing",)))]
100 #[pyo3(signature = (*_args, **_kwargs))]
101 fn py_new(_args: &Bound<'_, PyTuple>, _kwargs: Option<&Bound<'_, PyDict>>) -> Self {
102 Self
103 }
104
105 fn pre_process(&self, _data: &Bound<'_, PyAny>) {}
106
107 fn process(
108 &self,
109 _ts_now: u64,
110 _context: &PySimulationModuleContext,
111 ) -> PyResult<Option<Vec<Money>>> {
112 Err(to_pynotimplemented_err(
113 "Method 'process' must be implemented in a subclass.",
114 ))
115 }
116
117 fn acknowledge(&self, _outcomes: &Bound<'_, PyAny>) {}
118
119 fn log_diagnostics(&self) {}
120
121 fn reset(&self) {}
122}
123
124#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.backtest")]
126#[pyclass(
127 module = "nautilus_trader.backtest",
128 name = "SimulationModuleContext",
129 frozen,
130 unsendable
131)]
132#[derive(Debug)]
133pub struct PySimulationModuleContext {
134 venue: Venue,
135 base_currency: Option<Currency>,
136 instruments: Vec<InstrumentAny>,
137 order_books: Vec<OrderBook>,
138 positions: Vec<Position>,
139}
140
141impl PySimulationModuleContext {
142 fn from_exchange(ctx: &ExchangeContext<'_>) -> Self {
143 let mut instruments = ctx.instruments.values().cloned().collect::<Vec<_>>();
144 instruments.sort_unstable_by_key(Instrument::id);
145
146 let mut order_books = ctx
147 .matching_engines
148 .values()
149 .map(|engine| engine.get_book().clone())
150 .collect::<Vec<_>>();
151 order_books.sort_unstable_by_key(|book| book.instrument_id);
152
153 let mut positions = ctx
154 .cache
155 .positions_open(Some(&ctx.venue), None, None, None, None)
156 .into_iter()
157 .map(|position| (*position).clone())
158 .collect::<Vec<_>>();
159 positions.sort_unstable_by_key(|position| position.id);
160
161 Self {
162 venue: ctx.venue,
163 base_currency: ctx.base_currency,
164 instruments,
165 order_books,
166 positions,
167 }
168 }
169}
170
171#[pyo3_stub_gen::derive::gen_stub_pymethods]
172#[pymethods]
173impl PySimulationModuleContext {
174 #[getter]
175 fn venue(&self) -> Venue {
176 self.venue
177 }
178
179 #[getter]
180 fn base_currency(&self) -> Option<Currency> {
181 self.base_currency
182 }
183
184 #[getter]
185 fn instruments(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
186 self.instruments
187 .iter()
188 .cloned()
189 .map(|instrument| instrument_any_to_pyobject(py, instrument))
190 .collect()
191 }
192
193 #[getter]
194 fn order_books(&self) -> Vec<OrderBook> {
195 self.order_books.clone()
196 }
197
198 #[getter]
199 fn positions(&self) -> Vec<Position> {
200 self.positions.clone()
201 }
202}
203
204#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.backtest")]
206#[pyclass(
207 module = "nautilus_trader.backtest",
208 name = "AccountAdjustmentOutcome",
209 frozen
210)]
211#[derive(Debug, Clone)]
212pub struct PyAccountAdjustmentOutcome {
213 applied: bool,
214 error: Option<String>,
215}
216
217impl From<&AccountAdjustmentOutcome> for PyAccountAdjustmentOutcome {
218 fn from(outcome: &AccountAdjustmentOutcome) -> Self {
219 match outcome {
220 AccountAdjustmentOutcome::Applied => Self {
221 applied: true,
222 error: None,
223 },
224 AccountAdjustmentOutcome::Failed(error) => Self {
225 applied: false,
226 error: Some(error.to_string()),
227 },
228 }
229 }
230}
231
232#[pyo3_stub_gen::derive::gen_stub_pymethods]
233#[pymethods]
234impl PyAccountAdjustmentOutcome {
235 #[getter]
236 const fn applied(&self) -> bool {
237 self.applied
238 }
239
240 #[getter]
241 fn error(&self) -> Option<String> {
242 self.error.clone()
243 }
244}
245
246#[derive(Debug)]
247pub struct PythonSimulationModule {
248 obj: Py<PyAny>,
249}
250
251impl Clone for PythonSimulationModule {
252 fn clone(&self) -> Self {
253 Self::new(clone_py_object(&self.obj))
254 }
255}
256
257impl PythonSimulationModule {
258 #[must_use]
259 pub const fn new(obj: Py<PyAny>) -> Self {
260 Self { obj }
261 }
262
263 pub(crate) fn clone_ref(&self, py: Python<'_>) -> Py<PyAny> {
264 self.obj.clone_ref(py)
265 }
266}
267
268impl SimulationModule for PythonSimulationModule {
269 fn pre_process(&self, data: &Data) -> anyhow::Result<()> {
270 Python::attach(|py| -> anyhow::Result<()> {
271 let data = data_to_pyobject(py, data.clone())?;
272 self.obj.bind(py).call_method1("pre_process", (data,))?;
273 Ok(())
274 })
275 .map_err(|e| anyhow::anyhow!("Python SimulationModule.pre_process failed: {e}"))
276 }
277
278 fn process(
279 &self,
280 ts_now: nautilus_core::UnixNanos,
281 ctx: &ExchangeContext,
282 ) -> anyhow::Result<SimulationModuleResult> {
283 Python::attach(|py| -> anyhow::Result<SimulationModuleResult> {
284 let context = Py::new(py, PySimulationModuleContext::from_exchange(ctx))?;
285 let adjustments = self
286 .obj
287 .bind(py)
288 .call_method1("process", (ts_now.as_u64(), context))?
289 .extract::<Option<Vec<Money>>>()?;
290 Ok(adjustments.map_or(
291 SimulationModuleResult::NotReady,
292 SimulationModuleResult::Completed,
293 ))
294 })
295 .map_err(|e| anyhow::anyhow!("Python SimulationModule.process failed: {e}"))
296 }
297
298 fn acknowledge(&self, outcomes: &[AccountAdjustmentOutcome]) -> anyhow::Result<()> {
299 Python::attach(|py| -> anyhow::Result<()> {
300 let outcomes = outcomes
301 .iter()
302 .map(PyAccountAdjustmentOutcome::from)
303 .collect::<Vec<_>>();
304 self.obj.bind(py).call_method1("acknowledge", (outcomes,))?;
305 Ok(())
306 })
307 .map_err(|e| anyhow::anyhow!("Python SimulationModule.acknowledge failed: {e}"))
308 }
309
310 fn log_diagnostics(&self) -> anyhow::Result<()> {
311 Python::attach(|py| -> anyhow::Result<()> {
312 self.obj.bind(py).call_method0("log_diagnostics")?;
313 Ok(())
314 })
315 .map_err(|e| anyhow::anyhow!("Python SimulationModule.log_diagnostics failed: {e}"))
316 }
317
318 fn reset(&self) -> anyhow::Result<()> {
319 Python::attach(|py| -> anyhow::Result<()> {
320 self.obj.bind(py).call_method0("reset")?;
321 Ok(())
322 })
323 .map_err(|e| anyhow::anyhow!("Python SimulationModule.reset failed: {e}"))
324 }
325}
326
327fn pyobject_to_builtin_simulation_module_any(
328 obj: &Bound<'_, PyAny>,
329) -> Option<SimulationModuleAny> {
330 if let Ok(module) = obj.extract::<PyRef<'_, CfdSwapModule>>() {
331 return Some(SimulationModuleAny::CfdSwap((*module).clone()));
332 }
333
334 if let Ok(module) = obj.extract::<PyRef<'_, FXRolloverInterestModule>>() {
335 return Some(SimulationModuleAny::FXRolloverInterest((*module).clone()));
336 }
337 None
338}
339
340pub fn pyobject_to_simulation_module_any(obj: &Bound<'_, PyAny>) -> PyResult<SimulationModuleAny> {
346 if let Some(module) = pyobject_to_builtin_simulation_module_any(obj) {
347 return Ok(module);
348 }
349
350 if obj.is_instance_of::<PySimulationModule>() {
351 return Ok(SimulationModuleAny::Python(PythonSimulationModule::new(
352 obj.clone().unbind(),
353 )));
354 }
355
356 let type_name = obj.get_type().name()?;
357 Err(to_pytype_err(format!(
358 "Cannot convert {type_name} to SimulationModule"
359 )))
360}
361
362pub fn pyobject_to_simulation_module_handle(
370 py: Python<'_>,
371 obj: &Bound<'_, PyAny>,
372) -> PyResult<SimulationModuleHandle> {
373 if let Some(module) = pyobject_to_builtin_simulation_module_any(obj) {
374 return Ok(module.into());
375 }
376
377 let type_object = obj.get_type();
378 let type_id = type_object.as_ptr() as usize;
379 let extractor = SIMULATION_MODULE_EXTRACTORS.lock().get(&type_id).copied();
380
381 if let Some(extractor) = extractor {
382 return extractor(py, obj);
383 }
384
385 if obj.is_instance_of::<PySimulationModule>() {
386 return Ok(SimulationModuleHandle::new(PythonSimulationModule::new(
387 obj.clone().unbind(),
388 )));
389 }
390
391 Err(to_pytype_err(format!(
392 "Cannot convert {} to SimulationModule",
393 type_object.name()?
394 )))
395}
396
397pub fn simulation_module_any_to_pyobject(
403 py: Python<'_>,
404 module: &SimulationModuleAny,
405) -> PyResult<Py<PyAny>> {
406 match module {
407 SimulationModuleAny::CfdSwap(module) => Ok(Py::new(
408 py,
409 PyClassInitializer::from(PySimulationModule).add_subclass(module.clone()),
410 )?
411 .into_any()),
412 SimulationModuleAny::FXRolloverInterest(module) => Ok(Py::new(
413 py,
414 PyClassInitializer::from(PySimulationModule).add_subclass(module.clone()),
415 )?
416 .into_any()),
417 SimulationModuleAny::Python(module) => Ok(module.clone_ref(py)),
418 }
419}
420
421#[pyo3_stub_gen::derive::gen_stub_pymethods]
422#[pymethods]
423impl InterestRateRecord {
424 #[new]
426 fn py_new(location: String, time: String, value: Decimal) -> Self {
427 Self {
428 location,
429 time,
430 value,
431 }
432 }
433
434 fn __repr__(&self) -> String {
435 format!("{self:?}")
436 }
437}
438
439#[pyo3_stub_gen::derive::gen_stub_pymethods]
440#[pymethods]
441impl FXRolloverInterestModule {
442 #[new]
448 #[gen_stub(override_return_type(type_repr = "typing.Self", imports = ("typing",)))]
449 fn py_new(records: Vec<InterestRateRecord>) -> PyClassInitializer<Self> {
450 let module = Self::new(records);
451 PyClassInitializer::from(PySimulationModule).add_subclass(module)
452 }
453
454 fn __repr__(&self) -> String {
455 format!("{self:?}")
456 }
457}
458
459#[pyo3_stub_gen::derive::gen_stub_pymethods]
460#[pymethods]
461impl CfdSwapRate {
462 #[new]
464 fn py_new(instrument_id: InstrumentId, long_rate: Decimal, short_rate: Decimal) -> Self {
465 Self::new(instrument_id, long_rate, short_rate)
466 }
467
468 #[getter]
469 fn instrument_id(&self) -> InstrumentId {
470 self.instrument_id
471 }
472
473 #[getter]
474 fn long_rate(&self) -> Decimal {
475 self.long_rate
476 }
477
478 #[getter]
479 fn short_rate(&self) -> Decimal {
480 self.short_rate
481 }
482
483 fn __repr__(&self) -> String {
484 format!("{self:?}")
485 }
486}
487
488#[pyo3_stub_gen::derive::gen_stub_pymethods]
489#[pymethods]
490impl CfdSwapModule {
491 #[new]
493 #[gen_stub(override_return_type(type_repr = "typing.Self", imports = ("typing",)))]
494 #[pyo3(signature = (rates, rollover_hour=17, rollover_minute=0, triple_roll_weekday=5))]
495 fn py_new(
496 rates: Vec<CfdSwapRate>,
497 rollover_hour: i8,
498 rollover_minute: i8,
499 triple_roll_weekday: i8,
500 ) -> PyResult<PyClassInitializer<Self>> {
501 let rollover_time =
502 Time::new(rollover_hour, rollover_minute, 0, 0).map_err(to_pyvalue_err)?;
503 let triple_roll_weekday =
504 Weekday::from_monday_one_offset(triple_roll_weekday).map_err(to_pyvalue_err)?;
505 Ok(
506 PyClassInitializer::from(PySimulationModule).add_subclass(Self::new(
507 rates,
508 rollover_time,
509 triple_roll_weekday,
510 )),
511 )
512 }
513
514 fn __repr__(&self) -> String {
515 format!("{self:?}")
516 }
517}
518
519#[cfg(test)]
520mod tests {
521 use std::{cell::Cell, rc::Rc};
522
523 use indexmap::IndexMap;
524 use nautilus_common::cache::Cache;
525 use nautilus_execution::models::fee::{FeeModelAny, MakerTakerFeeModel};
526 use nautilus_model::{
527 data::Data,
528 enums::{AccountType, BookType, OmsType},
529 identifiers::Venue,
530 types::{Currency, Money},
531 };
532 use pyo3::{IntoPyObjectExt, exceptions::PyAttributeError, ffi::c_str, types::PyDict};
533 use rstest::rstest;
534
535 use super::*;
536 use crate::{
537 config::{BacktestEngineConfig, SimulatedVenueConfig},
538 engine::BacktestEngine,
539 };
540
541 fn with_empty_context<T>(f: impl FnOnce(&ExchangeContext<'_>) -> T) -> T {
542 let instruments = AHashMap::new();
543 let matching_engines = IndexMap::new();
544 let cache = Cache::default();
545 f(&ExchangeContext {
546 venue: Venue::new("SIM"),
547 base_currency: Some(Currency::USD()),
548 instruments: &instruments,
549 matching_engines: &matching_engines,
550 cache: &cache,
551 })
552 }
553
554 #[rstest]
555 fn test_pure_python_simulation_module_dispatch() {
556 Python::initialize();
557
558 Python::attach(|py| {
559 let locals = PyDict::new(py);
560 locals
561 .set_item("SimulationModule", py.get_type::<PySimulationModule>())
562 .unwrap();
563 let module = py
564 .eval(
565 c_str!(
566 "type('PurePythonSimulationModule', (SimulationModule,), {\
567 'process': lambda self, ts_now, context: \
568 (setattr(self, 'context', context), [self.adjustment])[1]\
569 })()"
570 ),
571 None,
572 Some(&locals),
573 )
574 .unwrap();
575 module
576 .setattr(
577 "adjustment",
578 Money::from("1.25 USD").into_py_any(py).unwrap(),
579 )
580 .unwrap();
581
582 assert!(matches!(
583 pyobject_to_simulation_module_any(&module).unwrap(),
584 SimulationModuleAny::Python(_)
585 ));
586 let handle = pyobject_to_simulation_module_handle(py, &module).unwrap();
587 let result =
588 with_empty_context(|ctx| handle.process(nautilus_core::UnixNanos::from(10), ctx))
589 .unwrap();
590 handle
591 .acknowledge(&[AccountAdjustmentOutcome::Applied])
592 .unwrap();
593 handle.reset().unwrap();
594
595 let context = module.getattr("context").unwrap();
596
597 assert_eq!(
598 result,
599 SimulationModuleResult::Completed(vec![Money::from("1.25 USD")])
600 );
601 assert_eq!(
602 context
603 .getattr("venue")
604 .unwrap()
605 .extract::<Venue>()
606 .unwrap(),
607 Venue::new("SIM")
608 );
609 assert_eq!(
610 context
611 .getattr("base_currency")
612 .unwrap()
613 .extract::<Option<Currency>>()
614 .unwrap(),
615 Some(Currency::USD())
616 );
617 assert_eq!(context.getattr("instruments").unwrap().len().unwrap(), 0);
618 assert_eq!(context.getattr("order_books").unwrap().len().unwrap(), 0);
619 assert_eq!(context.getattr("positions").unwrap().len().unwrap(), 0);
620 let error = context.setattr("venue", Venue::new("OTHER")).unwrap_err();
621 assert!(error.is_instance_of::<PyAttributeError>(py));
622 });
623 }
624
625 #[derive(Debug)]
626 struct NativeRustSimulationModule {
627 calls: Rc<Cell<u32>>,
628 }
629
630 impl SimulationModule for NativeRustSimulationModule {
631 fn pre_process(&self, _data: &Data) -> anyhow::Result<()> {
632 Ok(())
633 }
634
635 fn process(
636 &self,
637 _ts_now: nautilus_core::UnixNanos,
638 _ctx: &ExchangeContext,
639 ) -> anyhow::Result<SimulationModuleResult> {
640 self.calls.set(self.calls.get() + 1);
641 Ok(SimulationModuleResult::NotReady)
642 }
643
644 fn acknowledge(&self, _outcomes: &[AccountAdjustmentOutcome]) -> anyhow::Result<()> {
645 Ok(())
646 }
647
648 fn log_diagnostics(&self) -> anyhow::Result<()> {
649 Ok(())
650 }
651
652 fn reset(&self) -> anyhow::Result<()> {
653 Ok(())
654 }
655 }
656
657 #[pyclass(
658 name = "NativeSimulationModuleTest",
659 module = "native_simulation_module_test.registered",
660 unsendable
661 )]
662 #[derive(Debug)]
663 struct NativeSimulationModuleBinding {
664 calls: Rc<Cell<u32>>,
665 }
666
667 #[pyclass(
668 name = "NativeSimulationModuleTest",
669 module = "native_simulation_module_test.unregistered",
670 unsendable
671 )]
672 #[derive(Debug)]
673 struct UnregisteredNativeSimulationModuleBinding;
674
675 fn extract_native_simulation_module(
676 _py: Python<'_>,
677 obj: &Bound<'_, PyAny>,
678 ) -> PyResult<SimulationModuleHandle> {
679 let binding = obj.extract::<PyRef<'_, NativeSimulationModuleBinding>>()?;
680 Ok(SimulationModuleHandle::new(NativeRustSimulationModule {
681 calls: binding.calls.clone(),
682 }))
683 }
684
685 #[rstest]
686 fn test_registered_native_simulation_module_dispatch() {
687 Python::initialize();
688
689 Python::attach(|py| {
690 register_simulation_module_extractor::<NativeSimulationModuleBinding>(
691 py,
692 extract_native_simulation_module,
693 )
694 .unwrap();
695 let calls = Rc::new(Cell::new(0));
696 let binding = Py::new(
697 py,
698 NativeSimulationModuleBinding {
699 calls: calls.clone(),
700 },
701 )
702 .unwrap();
703 let handle =
704 pyobject_to_simulation_module_handle(py, binding.bind(py).as_any()).unwrap();
705
706 let result =
707 with_empty_context(|ctx| handle.process(nautilus_core::UnixNanos::from(10), ctx))
708 .unwrap();
709
710 assert_eq!(result, SimulationModuleResult::NotReady);
711 assert_eq!(calls.get(), 1);
712 });
713 }
714
715 #[rstest]
716 fn test_registered_native_simulation_module_uses_exact_python_type() {
717 Python::initialize();
718
719 Python::attach(|py| {
720 register_simulation_module_extractor::<NativeSimulationModuleBinding>(
721 py,
722 extract_native_simulation_module,
723 )
724 .unwrap();
725 let binding = Py::new(py, UnregisteredNativeSimulationModuleBinding).unwrap();
726
727 let error =
728 pyobject_to_simulation_module_handle(py, binding.bind(py).as_any()).unwrap_err();
729
730 assert_eq!(
731 error.to_string(),
732 "TypeError: Cannot convert NativeSimulationModuleTest to SimulationModule"
733 );
734 });
735 }
736
737 fn engine_with_module(module: SimulationModuleHandle) -> BacktestEngine {
738 let mut engine = BacktestEngine::new(BacktestEngineConfig::default()).unwrap();
739 engine
740 .add_venue(
741 SimulatedVenueConfig::builder()
742 .venue(Venue::new("SIM"))
743 .oms_type(OmsType::Netting)
744 .account_type(AccountType::Margin)
745 .book_type(BookType::L1_MBP)
746 .starting_balances(vec![Money::from("1000 USD")])
747 .modules(vec![module])
748 .fee_model(FeeModelAny::MakerTaker(MakerTakerFeeModel::zero()).into())
749 .build()
750 .unwrap(),
751 )
752 .unwrap();
753 engine
754 }
755
756 #[rstest]
757 fn test_python_simulation_module_exception_propagates_through_run() {
758 Python::initialize();
759
760 Python::attach(|py| {
761 let locals = PyDict::new(py);
762 locals
763 .set_item("SimulationModule", py.get_type::<PySimulationModule>())
764 .unwrap();
765 let module = py
766 .eval(
767 c_str!(
768 "type('FailingSimulationModule', (SimulationModule,), {\
769 'process': lambda self, ts_now, context: \
770 (_ for _ in ()).throw(ValueError('module boom'))\
771 })()"
772 ),
773 None,
774 Some(&locals),
775 )
776 .unwrap();
777 let handle = pyobject_to_simulation_module_handle(py, &module).unwrap();
778 let mut engine = engine_with_module(handle);
779
780 let error = engine.run(None, None, None, false).unwrap_err();
781 let message = error.to_string();
782 assert!(message.contains("Simulation module 0 process failed"));
783 assert!(message.contains("Python SimulationModule.process failed"));
784 assert!(message.contains("ValueError: module boom"));
785
786 assert_eq!(
787 engine.run(None, None, None, false).unwrap_err().to_string(),
788 format!("Simulation callback failure requires exchange reset: {message}")
789 );
790 });
791 }
792
793 #[rstest]
794 fn test_python_simulation_module_diagnostics_exception_propagates_through_run() {
795 Python::initialize();
796
797 Python::attach(|py| {
798 let locals = PyDict::new(py);
799 locals
800 .set_item("SimulationModule", py.get_type::<PySimulationModule>())
801 .unwrap();
802 let module = py
803 .eval(
804 c_str!(
805 "type('FailingDiagnosticsSimulationModule', (SimulationModule,), {\
806 'process': lambda self, ts_now, context: None,\
807 'log_diagnostics': lambda self: \
808 (_ for _ in ()).throw(ValueError('diagnostics boom'))\
809 })()"
810 ),
811 None,
812 Some(&locals),
813 )
814 .unwrap();
815 let handle = pyobject_to_simulation_module_handle(py, &module).unwrap();
816 let mut engine = engine_with_module(handle);
817
818 let error = engine.run(None, None, None, false).unwrap_err();
819 let message = error.to_string();
820
821 assert!(message.contains("Simulation module 0 log_diagnostics failed"));
822 assert!(message.contains("Python SimulationModule.log_diagnostics failed"));
823 assert!(message.contains("ValueError: diagnostics boom"));
824 });
825 }
826
827 #[rstest]
828 fn test_engine_reset_finishes_after_python_diagnostics_exception() {
829 Python::initialize();
830
831 Python::attach(|py| {
832 let locals = PyDict::new(py);
833 locals
834 .set_item("SimulationModule", py.get_type::<PySimulationModule>())
835 .unwrap();
836 let module = py
837 .eval(
838 c_str!(
839 "type('ResetAfterDiagnosticsFailureModule', (SimulationModule,), {\
840 'process': lambda self, ts_now, context: None,\
841 'log_diagnostics': lambda self: \
842 (_ for _ in ()).throw(ValueError('diagnostics boom')) \
843 if self.fail_diagnostics else None,\
844 'reset': lambda self: \
845 setattr(self, 'resets', self.resets + 1)\
846 })()"
847 ),
848 None,
849 Some(&locals),
850 )
851 .unwrap();
852 module.setattr("fail_diagnostics", true).unwrap();
853 module.setattr("resets", 0).unwrap();
854 let handle = pyobject_to_simulation_module_handle(py, &module).unwrap();
855 let mut engine = engine_with_module(handle);
856 engine.run(None, None, None, true).unwrap();
857
858 let error = engine.reset().unwrap_err();
859
860 assert!(
861 error
862 .to_string()
863 .contains("Simulation module 0 log_diagnostics failed")
864 );
865 assert_eq!(
866 module.getattr("resets").unwrap().extract::<u32>().unwrap(),
867 1
868 );
869 module.setattr("fail_diagnostics", false).unwrap();
870 engine.run(None, None, None, false).unwrap();
871 });
872 }
873}