Use EnvTestConfig for enabling invocation metering.

soroban-sdk/src/env.rs

@@ -254,15 +254,28 @@ struct EnvTestState {
 pub struct EnvTestConfig {
     /// Capture a test snapshot when the Env is dropped, causing a test snapshot
     /// JSON file to be written to disk when the Env is no longer referenced.
-    /// Defaults to true.
+    /// Defaults to `true``.
     pub capture_snapshot_at_drop: bool,
+    /// When `true` enables detailed per-invocation resource and fee estimation,
+    /// accessible via `env.cost_estimate()`.
+    ///
+    /// The top-level contract invocations and lifecycle operations (such as
+    /// `register` or `env.deployer()` operations) will be metered and
+    /// all the metering information will reset in-between them (note, that this
+    /// includes the CPU budget). Metering will not be reset while inside the
+    /// call (e.g. if a contract calls or creates another contract, that won't
+    /// reset metering).
+    ///     
+    /// Defaults to `false`.
+    pub enable_invocation_cost_estimate: bool,
 }
 
 #[cfg(any(test, feature = "testutils"))]
 impl Default for EnvTestConfig {
     fn default() -> Self {
         Self {
             capture_snapshot_at_drop: true,
+            enable_invocation_cost_estimate: false,
         }
     }
 }
@@ -490,6 +503,12 @@ impl Env {
         f((*self.test_state.generators).borrow_mut())
     }
 
+    fn apply_test_config(&self) {
+        if self.test_state.config.enable_invocation_cost_estimate {
+            self.env_impl.enable_invocation_metering();
+        }
+    }
+
     /// Create an Env with the test config.
     pub fn new_with_config(config: EnvTestConfig) -> Env {
         struct EmptySnapshotSource();
@@ -521,7 +540,13 @@ impl Env {
 
     /// Change the test config of an Env.
     pub fn set_config(&mut self, config: EnvTestConfig) {
+        if self.test_state.config.enable_invocation_cost_estimate
+            && !config.enable_invocation_cost_estimate
+        {
+            panic!("`enable_invocation_cost_estimate` flag may not be switched from `true` to `false`, a new `Env` instance has to be created instead.")
+        }
         self.test_state.config = config;
+        self.apply_test_config();
     }
 
     /// Used by multiple constructors to configure test environments consistently.
@@ -574,6 +599,7 @@ impl Env {
         };
 
         env.ledger().set(ledger_info);
+        env.apply_test_config();
 
         env
     }
@@ -593,7 +619,12 @@ impl Env {
     /// The returned value is as useful as the preceding setup, e.g. if a test
     /// contract is used instead of a Wasm contract, all the costs related to
     /// VM instantiation and execution, as well as Wasm reads/rent bumps will be
-    /// missed.    
+    /// missed.
+    ///
+    /// While the resource metering may be useful for contract optimization,
+    /// keep in mind that resource and fee estimation may be imprecise. Use
+    /// simulation with RPC in order to get the exact resources for submitting
+    /// the transactions to the network.
     pub fn cost_estimate(&self) -> CostEstimate {
         CostEstimate::new(self.clone())
     }
@@ -1629,7 +1660,7 @@ impl Env {
     }
 
     /// Get the budget that tracks the resources consumed for the environment.
-    #[deprecated(note = "use cost_estimate().detailed_metering()")]
+    #[deprecated(note = "use cost_estimate().budget()")]
     pub fn budget(&self) -> Budget {
         Budget::new(self.env_impl.budget_cloned())
     }

soroban-sdk/src/tests/budget.rs

@@ -21,12 +21,14 @@ fn test_budget() {
     let contract_id = e.register(Contract, ());
     let client = ContractClient::new(&e, &contract_id);
 
-    e.budget().reset_default();
+    e.cost_estimate().budget().reset_default();
     let b = client.add();
-    e.budget().print();
+    e.cost_estimate().budget().print();
 
     assert_eq!(
-        e.budget().tracker(ContractCostType::VisitObject),
+        e.cost_estimate()
+            .budget()
+            .tracker(ContractCostType::VisitObject),
         CostTracker {
             iterations: 13,
             inputs: None,

soroban-sdk/src/tests/cost_estimate.rs

@@ -1,4 +1,4 @@
-use crate as soroban_sdk;
+use crate::{self as soroban_sdk, env::EnvTestConfig};
 use expect_test::expect;
 use soroban_sdk::Env;
 use soroban_sdk_macros::symbol_short;
@@ -11,8 +11,10 @@ mod contract_data {
 // Update the test data in this test via running it with `UPDATE_EXPECT=1`.
 #[test]
 fn test_cost_estimate_with_storage() {
-    let e = Env::default();
-    e.cost_estimate().enable();
+    let e = Env::new_with_config(EnvTestConfig {
+        enable_invocation_cost_estimate: true,
+        ..Default::default()
+    });
 
     let contract_id = e.register(contract_data::WASM, ());
     let client = contract_data::Client::new(&e, &contract_id);

soroban-sdk/src/tests/env.rs

@@ -166,6 +166,7 @@ fn test_snapshot_file_disabled() {
         let _ = e1.register(Contract, ());
         let e2 = Env::new_with_config(EnvTestConfig {
             capture_snapshot_at_drop: false,
+            ..Default::default()
         });
         let _ = e2.register(Contract, ());
         assert!(!p1.exists());
@@ -194,6 +195,7 @@ fn test_snapshot_file_disabled_after_creation() {
         let mut e2 = Env::default();
         e2.set_config(EnvTestConfig {
             capture_snapshot_at_drop: false,
+            ..Default::default()
         });
         let _ = e2.register(Contract, ());
         assert!(!p1.exists());

soroban-sdk/src/tests/proptest_scval_cmp.rs

@@ -70,7 +70,7 @@ proptest! {
         prop_assert_eq!(scval_partial_eq, rawval_cmp_is_eq);
 
         // Compare<ScVal> for Budget
-        let budget = env.budget().0;
+        let budget = env.cost_estimate().budget().0;
         let scval_budget_cmp = budget.compare(&scval_1, &scval_2).expect("cmp");
 
         if scval_budget_cmp != scval_cmp {

soroban-sdk/src/tests/proptest_val_cmp.rs

@@ -20,7 +20,7 @@ proptest! {
         rawval_proto_2 in arb::<<Val as SorobanArbitrary>::Prototype>(),
     ) {
         let env = &Env::default();
-        let budget = env.budget().0;
+        let budget = env.cost_estimate().budget().0;
 
         let rawval_1 = Val::from_val(env, &rawval_proto_1);
         let rawval_2 = Val::from_val(env, &rawval_proto_2);

soroban-sdk/src/testutils.rs

@@ -314,9 +314,9 @@ pub mod budget {
     /// # #[cfg(feature = "testutils")]
     /// # fn main() {
     /// #     let env = Env::default();
-    /// env.budget().reset_default();
+    /// env.cost_estimate().budget().reset_default();
     /// // ...
-    /// println!("{}", env.budget());
+    /// println!("{}", env.cost_estimate().budget());
     /// # }
     /// # #[cfg(not(feature = "testutils"))]
     /// # fn main() { }

soroban-sdk/src/testutils/cost_estimate.rs

@@ -11,27 +11,6 @@ impl CostEstimate {
         Self { env }
     }
 
-    /// Enables detailed per-invocation resource cost metering.
-    ///
-    /// The top-level contract invocations and lifecycle operations (such as
-    /// `register` or `env.deployer()` operations) will be metered and
-    /// all the metering information will reset in-between them. Metering will
-    /// not be reset while inside the call (e.g. if a contract calls or creates
-    /// another contract, that won't reset metering).
-    ///
-    /// The metered resources for the last invocation can be retrieved with
-    /// `resources()`, the estimated fee corresponding to these resources can be
-    /// retrieved with `fee()`, and the detailed CPU and memory memory
-    /// breakdown can be retrieved with `detailed_metering()`.
-    ///
-    /// While the resource metering may be useful for contract optimization,
-    /// keep in mind that resource and fee estimation may be imprecise. Use
-    /// simulation with RPC in order to get the exact resources for submitting
-    /// the transactions to the network.    
-    pub fn enable(&self) {
-        self.env.host().enable_invocation_metering();
-    }
-
     /// Returns the resources metered during the last top level contract
     /// invocation.    
     /// Take the return value with a grain of salt. The returned resources mostly
@@ -90,16 +69,20 @@ impl CostEstimate {
         )
     }
 
-    /// Returns the detailed CPU and memory metering information recorded thus
-    /// far.
+    /// Returns the budget object that provides the detailed CPU and memory
+    /// metering information recorded thus far.
     ///
-    /// The metering resets before every top-level contract level invocation.
+    /// The budget metering resets before every top-level contract level
+    /// invocation.
+    /// 
+    /// budget() may also be used to adjust the CPU and memory limits via the
+    /// `reset_` methods.
     ///
     /// Note, that unlike `resources()`/`fee()` this will always return some
     /// value. If there was no contract call, then the resulting value will
     /// correspond to metering any environment setup that has been made thus
     /// far.
-    pub fn detailed_metering(&self) -> Budget {
+    pub fn budget(&self) -> Budget {
         Budget::new(self.env.host().budget_cloned())
     }
 }