Use checked_feature

near · Apr 6, 2023 · eb03ed5 · eb03ed5
1 parent e9f6eed
commit eb03ed5
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/runtime/runtime/src/lib.rs b/runtime/runtime/src/lib.rs
@@ -1287,9 +1287,7 @@ impl Runtime {
                     .expect("`process_transaction` must populate compute usage"),
             )?;
 
-            if apply_state.current_protocol_version
-                < ProtocolFeature::ComputeCosts.protocol_version()
-            {
+            if !checked_feature!("stable", ComputeCosts, apply_state.current_protocol_version) {
                 assert_eq!(
                     total_compute_usage, total_gas_burnt,
                     "Compute usage must match burnt gas"
@@ -1339,9 +1337,7 @@ impl Runtime {
                         .expect("`process_receipt` must populate compute usage"),
                 )?;
 
-                if apply_state.current_protocol_version
-                    < ProtocolFeature::ComputeCosts.protocol_version()
-                {
+                if !checked_feature!("stable", ComputeCosts, apply_state.current_protocol_version) {
                     assert_eq!(
                         total_compute_usage, total_gas_burnt,
                         "Compute usage must match burnt gas"
@@ -2562,8 +2558,18 @@ mod tests {
             setup_runtime(initial_balance, initial_locked, 1);
 
         let mut free_config = RuntimeConfig::free();
-        let sha256_cost =
-            ParameterCost { gas: Gas::from(1u64), compute: Compute::from(10_000_000_000_000u64) };
+        let sha256_cost = ParameterCost {
+            gas: Gas::from(1_000_000u64),
+            compute: if checked_feature!(
+                "stable",
+                ComputeCosts,
+                apply_state.current_protocol_version
+            ) {
+                Compute::from(10_000_000_000_000u64)
+            } else {
+                Compute::from(1_000_000u64)
+            },
+        };
         free_config.wasm_config.ext_costs.costs[ExtCosts::sha256_base] = sha256_cost.clone();
         apply_state.config = Arc::new(free_config);
         // This allows us to execute 1 receipt with a function call per apply.
@@ -2583,7 +2589,7 @@ mod tests {
             vec![Action::FunctionCall(FunctionCallAction {
                 method_name: "ext_sha256".to_string(),
                 args: b"first".to_vec(),
-                gas: 1,
+                gas: sha256_cost.gas,
                 deposit: 0,
             })],
         );
@@ -2594,7 +2600,7 @@ mod tests {
             vec![Action::FunctionCall(FunctionCallAction {
                 method_name: "ext_sha256".to_string(),
                 args: b"second".to_vec(),
-                gas: 1,
+                gas: sha256_cost.gas,
                 deposit: 0,
             })],
         );