1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
//! This module defines and implements the [`TestEnvironment`] struct.

use super::*;
use crate::prelude::*;

/// The environment that all tests of this testing framework are run against.
///
/// This struct may be thought of as the main struct in this testing framework which encapsulates a
/// a self-contained instance of the Radix Engine ([`EncapsulatedRadixEngine`]). The functionality
/// of the Radix Engine is exposed through the [`ClientApi`] which makes this testing environment no
/// less capable than Scrypto code.
///
/// ## Introduction
///
/// This testing framework is designed to allow you to write Scrypto-like code and use that to test
/// your packages and blueprints and follows a different approach from the `LedgerSimulator` class. The
/// test-runner is an in-memory ledger simulator which you can interact with as a user that submits
/// transactions to the network. The approach followed by this testing framework is different,
/// instead of submitting transactions, you're making invocations to the Radix Engine, getting
/// results back, and then writing assertions against what you got back.
///
/// Both the LedgerSimulator and this testing framework will prove to be useful throughout your blueprint
/// development journey. As an example, this testing framework allows you to disable some of kernel
/// modules that may get in your way when writing tests so it may be an optimal framework to use to
/// ensure that the "math checks out" in your blueprint code without needing to think about costing
/// or auth. However, when you're reaching the final stages of developing a blueprint you may want
/// tests that check that interactions with your blueprint will succeed in a simulated setting that
/// is close to the real setting, which is when the `LedgerSimulator` comes in. Overall, we may put these
/// two frameworks into two categories: This framework (named scrypto-test) is a framework for unit
/// testing your blueprints and is a good framework to use to check that your DeFi logic is correct.
/// The `LedgerSimulator` is an integration testing or an end-to-end testing framework to test that your
/// blueprints work in a simulated ledger with all of the costing limits, substate limits, and other
/// limits applied.
///
/// ## Features
///
/// This framework has many new features that developers may find useful when testing their packages
/// some of those features are:
///
/// * The ability to create mock [`Bucket`]s and [`Proof`]s through two main ways: by creating them
///   out of thin air, and by disabling the auth module and minting them. This functionality can be
///   found in the [`BucketFactory`] and [`ProofFactory`] structs and the [`CreationStrategy`].
/// * The ability to query the contents of [`Bucket`]s and [`Proof`]s for the purpose of writing
///   assertions against them. Not only that, but this testing framework allows you to call any
///   method you wish on these nodes. As an example, in a test, you can get a [`Bucket`] and then
///   create a proof out of it in manner similar to Scrypto.
/// * The ability to enable and disable kernel modules at runtime. The Radix Engine kernel is quite
///   modular with concepts such as auth, costing, and limits being implemented as kernel modules.
///   Disabling or enabling kernel modules at runtime can prove to be quite useful when writing DeFi
///   tests. As an example, you may want to not think about costing at all when writing tests and
///   thus you may opt to disable the costing module entirely and continue your test without it.
///   This can be done through [`TestEnvironment::disable_costing_module`].
/// * This testing framework uses test bindings to provide a higher-level API for calling methods
///   and functions on a blueprint without the need to do raw [`TestEnvironment::call_method_typed`]
///   or [`TestEnvironment::call_function_typed`]. The test bindings are generated by the blueprint
///   macro and are feature gated behind a `test` feature.
///
/// ## Getting Started
///
/// The following example shows a very simple test that gets XRD from the faucet and then asserts
/// that the amount is equal to what we expect.
///
/// ```
/// use scrypto_test::prelude::*;
///
/// // Arrange
/// let mut env = TestEnvironment::new();
///
/// // Act
/// let bucket = env.call_method_typed::<_, _, Bucket>(FAUCET, "free", &()).unwrap();
///
/// // Assert
/// let amount = bucket.amount(&mut env).unwrap();
/// assert_eq!(amount, dec!("10000"));
/// ```
///
/// A few things to note about the code you see above:
///
/// * There is no transactions, worktop, receipt, manifests or anything of that sort! This part is
///   "not just hidden" from this testing framework but is actually non existent! The approach that
///   framework of wrapping a self-contained Radix Engine means that there is no need for manifests
///   or other transaction concepts.
/// * Methods such as [`Bucket::amount`] can be called to get the amount of resources in a bucket
///   and then assert against that.
///
/// ## Manipulating Kernel Modules
///
/// At runtime, the kernel modules can be enabled or disabled. For each kernel module there are four
/// methods on the [`TestEnvironment`]:
///
/// * A method to enable the kernel module (e.g., [`TestEnvironment::enable_costing_module`]).
/// * A method to disable the kernel module (e.g., [`TestEnvironment::disable_costing_module`]).
/// * A method to perform some action in a callback with the module enabled (e.g.,
/// [`TestEnvironment::with_costing_module_enabled`]).
/// * A method to perform some action in a callback with the module disabled (e.g.,
/// [`TestEnvironment::with_costing_module_disabled`]).
///
/// The simple enable and disable methods are quite straightforward: call them to enable or disable
/// a kernel module. The `with_*` methods are a little bit more intricate, they allow you to perform
/// some actions with a specific kernel either enabled or disabled and then resets the state of the
/// kernel modules afterwards. As an example:
///
/// ```
/// use scrypto_test::prelude::*;
///
/// // Arrange
/// let mut env = TestEnvironment::new();
///
/// // Act
/// let bucket = env.with_auth_module_disabled(|env| {
///     /* Auth Module is disabled just before this point */
///     ResourceManager(XRD).mint_fungible(100.into(), env).unwrap()
///     /* Kernel modules are reset just after this point. */
/// });
///
/// // Assert
/// let amount = bucket.amount(&mut env).unwrap();
/// assert_eq!(amount, dec!("100"))
/// ```
///
/// ## Common Arranges or Teardowns
///
/// There are cases where you may have many tests that all share a large portion of your arrange
/// or teardown logic. While this framework does not specifically provide solutions for this, there
/// are many useful Rust patterns that may be employed here to allow you to do this: the simplest
/// and the most elegant is probably by using callback functions.
///
/// Imagine this, you're building a Dex and many of the tests you write require you to have two
/// resources with a very large supply so you can write your tests with. You can achieve this by
/// doing something like:
///
/// ```
/// use scrypto_test::prelude::*;
///
/// pub fn two_resource_environment<F>(func: F)
/// where
///     F: FnOnce(TestEnvironment, Bucket, Bucket),
/// {
///     let mut env = TestEnvironment::new();
///     let bucket1 = ResourceBuilder::new_fungible(OwnerRole::None)
///         .mint_initial_supply(dec!("100000000000"), &mut env)
///         .unwrap();
///     let bucket2 = ResourceBuilder::new_fungible(OwnerRole::None)
///         .mint_initial_supply(dec!("100000000000"), &mut env)
///         .unwrap();
///     func(env, bucket1, bucket2)
///
///     /* Potential teardown happens here */
/// }
///
/// #[test]
/// fn contribution_provides_expected_amount_of_pool_units() {
///     two_resource_environment(|mut env, bucket1, bucket2| {
///         /* Your test goes here */
///     })
/// }
/// ```
///
/// You may have a function like `two_resource_environment` seen above which sets up the environment
/// and then some callback and potentially then executes some teardown code. Another way to do this
/// would be through simple factory and destructor methods.
/// ```
pub struct TestEnvironment<D>(pub(super) EncapsulatedRadixEngine<D>)
where
    D: SubstateDatabase + CommittableSubstateDatabase + 'static;

impl TestEnvironment<InMemorySubstateDatabase> {
    pub fn new() -> Self {
        TestEnvironmentBuilder::new().build()
    }
}

impl<D> TestEnvironment<D>
where
    D: SubstateDatabase + CommittableSubstateDatabase + 'static,
{
    //=============
    // Invocations
    //=============

    /// Invokes a function on the provided blueprint and package with the given arguments.
    ///
    /// This method is a typed version of the [`ClientBlueprintApi::call_function`] which Scrypto
    /// encodes the arguments and Scrypto decodes the returns on behalf of the caller. This method
    /// assumes that the caller is correct about the argument and return types and panics if the
    /// encoding or decoding fails.
    ///
    /// # Arguments
    ///
    /// * `package_address`: [`PackageAddress`] - The address of the package that contains the
    /// blueprint.
    /// * `blueprint_name`: [`&str`] - The name of the blueprint.
    /// * `function_name`: [`&str`] - The nae of the function.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    /// is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    /// was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    /// [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    /// [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    /// the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    /// could be decoded as the desired output type.
    pub fn call_function_typed<I, O>(
        &mut self,
        package_address: PackageAddress,
        blueprint_name: &str,
        function_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_function(package_address, blueprint_name, function_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on the main module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`ClientObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    /// fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    /// is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    /// was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    /// [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    /// [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    /// the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    /// could be decoded as the desired output type.
    pub fn call_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_method(&node_id.into(), method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on the main module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`ClientObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    /// fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    /// is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    /// was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    /// [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    /// [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    /// the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    /// could be decoded as the desired output type.
    pub fn call_direct_access_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_direct_access_method(&node_id.into(), method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    /// Invokes a method on a module of a node with the provided typed arguments.
    ///
    /// This method is a typed version of the [`ClientObjectApi::call_method`] which Scrypto encodes
    /// the arguments and Scrypto decodes the returns on behalf of the caller. This method assumes
    /// that the caller is correct about the argument and return types and panics if the encoding or
    /// decoding fails.
    ///
    /// # Arguments
    ///
    /// * `node_id`: `T` - The node to invoke the method on. This is a generic argument that's
    /// fulfilled by any type that implements [`Into<NodeId>`], thus, any address type can be used.
    /// * `module`: [`AttachedModuleId`] - The module id.
    /// * `method_name`: [`&str`] - The name of the method to invoke.
    /// * `args`: `&I` - The arguments to invoke the method with. This is a generic arguments that
    /// is fulfilled by any type that implements [`ScryptoEncode`].
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The returns from the method invocation. If the invocation
    /// was successful a [`Result::Ok`] is returned, otherwise a [`Result::Err`] is returned. The
    /// [`Result::Ok`] variant is a generic that's fulfilled by any type that implements
    /// [`ScryptoDecode`].
    ///
    /// # Panics
    ///
    /// This method panics in the following two cases:
    ///
    /// * Through an unwrap when calling [`scrypto_encode`] on the method arguments. Please consult
    /// the SBOR documentation on more information on why SBOR encoding may fail.
    /// * Through an unwrap when calling [`scrypto_decode`] on the returns. This panics if the type
    /// could be decoded as the desired output type.
    pub fn call_module_method_typed<N, I, O>(
        &mut self,
        node_id: N,
        module: AttachedModuleId,
        method_name: &str,
        args: &I,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId>,
        I: ScryptoEncode,
        O: ScryptoDecode,
    {
        let args = scrypto_encode(args).expect("Scrypto encoding of args failed");
        self.call_module_method(&node_id.into(), module, method_name, args)
            .map(|rtn| scrypto_decode(&rtn).expect("Scrypto decoding of returns failed"))
    }

    //====================================
    // Manipulation of the Kernel Modules
    //====================================

    /// Enables the kernel trace kernel module of the Radix Engine.
    pub fn enable_kernel_trace_module(&mut self) {
        self.enable_module(EnabledModules::KERNEL_TRACE)
    }

    /// Enables the limits kernel module of the Radix Engine.
    pub fn enable_limits_module(&mut self) {
        self.enable_module(EnabledModules::LIMITS)
    }

    /// Enables the costing kernel module of the Radix Engine.
    pub fn enable_costing_module(&mut self) {
        self.enable_module(EnabledModules::COSTING)
    }

    /// Enables the auth kernel module of the Radix Engine.
    pub fn enable_auth_module(&mut self) {
        self.enable_module(EnabledModules::AUTH)
    }

    /// Enables the transaction env kernel module of the Radix Engine.
    pub fn enable_transaction_runtime_module(&mut self) {
        self.enable_module(EnabledModules::TRANSACTION_RUNTIME)
    }

    /// Enables the execution trace kernel module of the Radix Engine.
    pub fn enable_execution_trace_module(&mut self) {
        self.enable_module(EnabledModules::EXECUTION_TRACE)
    }

    /// Disables the kernel trace kernel module of the Radix Engine.
    pub fn disable_kernel_trace_module(&mut self) {
        self.disable_module(EnabledModules::KERNEL_TRACE)
    }

    /// Disables the limits kernel module of the Radix Engine.
    pub fn disable_limits_module(&mut self) {
        self.disable_module(EnabledModules::LIMITS)
    }

    /// Disables the costing kernel module of the Radix Engine.
    pub fn disable_costing_module(&mut self) {
        self.disable_module(EnabledModules::COSTING)
    }

    /// Disables the auth kernel module of the Radix Engine.
    pub fn disable_auth_module(&mut self) {
        self.disable_module(EnabledModules::AUTH)
    }

    /// Disables the transaction env kernel module of the Radix Engine.
    pub fn disable_transaction_runtime_module(&mut self) {
        self.disable_module(EnabledModules::TRANSACTION_RUNTIME)
    }

    /// Disables the execution trace kernel module of the Radix Engine.
    pub fn disable_execution_trace_module(&mut self) {
        self.disable_module(EnabledModules::EXECUTION_TRACE)
    }

    /// Calls the passed `callback` with the kernel trace kernel module enabled and then resets the
    /// state of the kernel modules.
    pub fn with_kernel_trace_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_kernel_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the limits kernel module enabled and then resets the state
    /// of the kernel modules.
    pub fn with_limits_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_limits_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the costing kernel module enabled and then resets the state
    /// of the kernel modules.
    pub fn with_costing_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_costing_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the auth kernel module enabled and then resets the state of
    /// the kernel modules.
    pub fn with_auth_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_auth_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the transaction env kernel module enabled and then
    /// resets the state of the kernel modules.
    pub fn with_transaction_runtime_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_transaction_runtime_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the execution trace kernel module enabled and then resets
    /// the state of the kernel modules.
    pub fn with_execution_trace_module_enabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.enable_execution_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the kernel trace kernel module disabled and then resets the
    /// state of the kernel modules.
    pub fn with_kernel_trace_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_kernel_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the limits kernel module disabled and then resets the state
    /// of the kernel modules.
    pub fn with_limits_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_limits_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the costing kernel module disabled and then resets the
    /// state of the kernel modules.
    pub fn with_costing_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_costing_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the auth kernel module disabled and then resets the state
    /// of the kernel modules.
    pub fn with_auth_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_auth_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the transaction env kernel module disabled and then
    /// resets the state of the kernel modules.
    pub fn with_transaction_runtime_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_transaction_runtime_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Calls the passed `callback` with the execution trace kernel module disabled and then resets
    /// the state of the kernel modules.
    pub fn with_execution_trace_module_disabled<F, O>(&mut self, callback: F) -> O
    where
        F: FnOnce(&mut Self) -> O,
    {
        let enabled_modules = self.enabled_modules();
        self.disable_execution_trace_module();
        let rtn = callback(self);
        self.set_enabled_modules(enabled_modules);
        rtn
    }

    /// Returns the bit flags representing the currently enabled kernel modules.
    pub fn enabled_modules(&self) -> EnabledModules {
        self.0
            .with_kernel(|kernel| kernel.kernel_callback().modules.enabled_modules)
    }

    /// Sets the bit flags representing the enabled kernel modules.
    pub fn set_enabled_modules(&mut self, enabled_modules: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules = enabled_modules
        })
    }

    /// Enables specific kernel module(s).
    pub fn enable_module(&mut self, module: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules |= module
        })
    }

    /// Disables specific kernel module(s).
    pub fn disable_module(&mut self, module: EnabledModules) {
        self.0.with_kernel_mut(|kernel| {
            kernel.kernel_callback_mut().modules.enabled_modules &= !module
        })
    }

    //=======
    // State
    //=======

    /// Reads the state of a component and allows for a callback to be executed over the decoded
    /// state.
    ///
    /// This method performs the steps needed to read the state of a component and then perform the
    /// various steps needed before the state can be read or used such as the locking, reading, and
    /// decoding of the substate and the various steps that need to be performed after the state is
    /// read such as unlocking the substate.
    ///
    /// Users of this method are expected to pass in a callback function to operate over the state
    /// as this is the main way that this method ensures that references do not escape out of this
    /// method after the substate is closed.
    ///
    /// # Arguments
    ///
    /// * `node_id`: [`N`] - The address of the component to read the state of. This is a generic
    /// type parameter that's satisfied by any type that implements [`Into<NodeId>`].
    /// * `callback`: [`F`] - A callback function to call after the component state has been read
    /// and decoded into the type specified by the generic parameter [`S`]. Anything returned from
    /// this callback is returned from this method unless an error happens after the callback is
    /// executed.
    ///
    /// # Returns
    ///
    /// * [`Result<O, RuntimeError>`] - The output of the callback function passed in or a runtime
    /// error if one of the steps failed.
    ///
    /// # Panics
    ///
    /// This method panics if the component state can not be decoded as the generic type parameter
    /// [`S`].
    pub fn with_component_state<S, N, F, O>(
        &mut self,
        node_id: N,
        mut callback: F,
    ) -> Result<O, RuntimeError>
    where
        S: ScryptoDecode,
        N: Into<NodeId>,
        F: FnMut(&mut S, &mut Self) -> O,
    {
        let (handle, state) = self.0.with_kernel_mut(|kernel| {
            // Lock
            let handle = kernel.kernel_open_substate(
                &node_id.into(),
                MAIN_BASE_PARTITION,
                &SubstateKey::Field(ComponentField::State0.into()),
                LockFlags::read_only(),
                SystemLockData::Field(FieldLockData::Read),
            )?;

            // Read
            let state = kernel.kernel_read_substate(handle).map(|v| {
                let FieldSubstate::<ScryptoValue>::V1(FieldSubstateV1 { payload, .. }) =
                    v.as_typed().unwrap();
                scrypto_encode(&payload).unwrap()
            })?;

            Ok::<_, RuntimeError>((handle, state))
        })?;

        // Decode
        let mut state = scrypto_decode::<S>(&state).unwrap();

        // Callback
        let rtn = callback(&mut state, self);

        // Unlock
        self.0
            .with_kernel_mut(|kernel| kernel.kernel_close_substate(handle))?;

        Ok(rtn)
    }

    //===================
    // Epoch & Timestamp
    //===================

    /// Gets the current epoch from the Consensus Manager.
    pub fn get_current_epoch(&mut self) -> Epoch {
        Runtime::current_epoch(self).unwrap()
    }

    /// Sets the current epoch.
    pub fn set_current_epoch(&mut self, epoch: Epoch) {
        self.as_method_actor(
            CONSENSUS_MANAGER,
            ModuleId::Main,
            CONSENSUS_MANAGER_NEXT_ROUND_IDENT,
            |env| -> Result<(), RuntimeError> {
                let manager_handle = env
                    .actor_open_field(
                        ACTOR_STATE_SELF,
                        ConsensusManagerField::State.into(),
                        LockFlags::MUTABLE,
                    )
                    .unwrap();
                let mut manager_substate =
                    env.field_read_typed::<VersionedConsensusManagerState>(manager_handle)?;

                manager_substate.as_unique_version_mut().epoch = epoch;

                env.field_write_typed(manager_handle, &manager_substate)?;
                env.field_close(manager_handle)?;
                Ok(())
            },
        )
        .unwrap()
        .unwrap();
    }

    /// Gets the current time stamp from the Consensus Manager.
    pub fn get_current_time(&mut self) -> Instant {
        Runtime::current_time(self, TimePrecision::Second).unwrap()
    }

    pub fn set_current_time(&mut self, instant: Instant) {
        self.as_method_actor(
            CONSENSUS_MANAGER,
            ModuleId::Main,
            CONSENSUS_MANAGER_NEXT_ROUND_IDENT,
            |env| -> Result<(), RuntimeError> {
                let handle = env.actor_open_field(
                    ACTOR_STATE_SELF,
                    ConsensusManagerField::ProposerMilliTimestamp.into(),
                    LockFlags::MUTABLE,
                )?;
                let mut proposer_milli_timestamp =
                    env.field_read_typed::<ConsensusManagerProposerMilliTimestampFieldPayload>(
                        handle,
                    )?;
                proposer_milli_timestamp.as_unique_version_mut().epoch_milli =
                    instant.seconds_since_unix_epoch * 1000;
                env.field_write_typed(handle, &proposer_milli_timestamp)?;
                env.field_close(handle)?;
                Ok(())
            },
        )
        .unwrap()
        .unwrap();
    }

    //=========
    // Helpers
    //=========

    /// Allows us to perform some action as another actor.
    ///
    /// This function pushes a new call-frame onto the stack with the actor information we desire,
    /// performs the call-back, and then pops the call-frame off.
    pub(crate) fn as_method_actor<N, F, O>(
        &mut self,
        node_id: N,
        module_id: ModuleId,
        method_ident: &str,
        callback: F,
    ) -> Result<O, RuntimeError>
    where
        N: Into<NodeId> + Copy,
        F: FnOnce(&mut Self) -> O,
        O: ScryptoEncode,
    {
        let object_info = self.0.with_kernel_mut(|kernel| {
            SystemService {
                api: kernel,
                phantom: PhantomData,
            }
            .get_object_info(&node_id.into())
        })?;
        let auth_zone = self.0.with_kernel_mut(|kernel| {
            let mut system_service = SystemService {
                api: kernel,
                phantom: PhantomData,
            };
            AuthModule::on_call_fn_mock(
                &mut system_service,
                Some((&node_id.into(), false)),
                Default::default(),
                Default::default(),
            )
        })?;
        let actor = Actor::Method(MethodActor {
            method_type: match module_id {
                ModuleId::Main => MethodType::Main,
                ModuleId::Royalty => MethodType::Module(AttachedModuleId::Royalty),
                ModuleId::Metadata => MethodType::Module(AttachedModuleId::Metadata),
                ModuleId::RoleAssignment => MethodType::Module(AttachedModuleId::RoleAssignment),
            },
            ident: method_ident.to_owned(),
            node_id: node_id.into(),
            auth_zone,
            object_info,
        });
        self.as_actor(actor, callback)
    }

    /// Allows us to perform some action as another actor.
    ///
    /// This function pushes a new call-frame onto the stack with the actor information we desire,
    /// performs the call-back, and then pops the call-frame off.
    pub(crate) fn as_actor<F, O>(&mut self, actor: Actor, callback: F) -> Result<O, RuntimeError>
    where
        F: FnOnce(&mut Self) -> O,
        O: ScryptoEncode,
    {
        // Creating the next frame.
        let mut message =
            CallFrameMessage::from_input(&IndexedScryptoValue::from_typed(&()), &actor);
        self.0.with_kernel_mut(|kernel| {
            let (substate_io, current_frame) = kernel.kernel_current_frame_mut();
            message
                .copy_global_references
                .extend(current_frame.stable_references().keys());
            let new_frame =
                CallFrame::new_child_from_parent(substate_io, current_frame, actor, message)
                    .unwrap();
            let old = core::mem::replace(current_frame, new_frame);
            kernel.kernel_prev_frame_stack_mut().push(old);
        });

        // Executing the callback
        let rtn = callback(self);

        // Constructing the message from the returns
        let message = {
            let indexed_scrypto_value = IndexedScryptoValue::from_typed(&rtn);
            CallFrameMessage {
                move_nodes: indexed_scrypto_value.owned_nodes().clone(),
                copy_global_references: indexed_scrypto_value.references().clone(),
                ..Default::default()
            }
        };

        // Popping the last frame & Passing message
        self.0
            .with_kernel_mut(|kernel| -> Result<(), RuntimeError> {
                let mut previous_frame = kernel.kernel_prev_frame_stack_mut().pop().unwrap();
                let (substate_io, current_frame) = kernel.kernel_current_frame_mut();

                CallFrame::pass_message(
                    substate_io,
                    current_frame,
                    &mut previous_frame,
                    message.clone(),
                )
                .map_err(CallFrameError::PassMessageError)
                .map_err(KernelError::CallFrameError)?;

                *current_frame = previous_frame;
                Ok(())
            })?;

        Ok(rtn)
    }
}

impl Default for TestEnvironment<InMemorySubstateDatabase> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use crate::prelude::*;

    #[test]
    pub fn test_env_can_be_created() {
        let _ = TestEnvironment::new();
    }
}