@layerzerolabs/protocol-stellar-v2 0.2.40 → 0.2.41

package/contracts/oapps/oft/src/tests/extensions/pausable.rs

@@ -1,7 +1,9 @@
 extern crate std;
 
 use crate::extensions::pausable::{OFTPausable, OFTPausableError, OFTPausableInternal};
-use soroban_sdk::{contract, contractimpl, Address, Env};
+use crate::extensions::pausable::{PAUSER_ROLE, UNPAUSER_ROLE};
+use soroban_sdk::{contract, contractimpl, Address, Env, Symbol};
+use utils::rbac::{grant_role_no_auth, RoleBasedAccessControl};
 use utils::auth::Auth;
 
 // ============================================================================
@@ -22,8 +24,18 @@ impl OFTPausableInternal for PausableTestContract {}
 #[contractimpl(contracttrait)]
 impl OFTPausable for PausableTestContract {}
 
+#[contractimpl(contracttrait)]
+impl RoleBasedAccessControl for PausableTestContract {}
+
 #[contractimpl]
 impl PausableTestContract {
+    /// Test-only: grants PAUSER_ROLE and UNPAUSER_ROLE to the contract.
+    pub fn init_roles(env: Env) {
+        let contract_id = env.current_contract_address();
+        grant_role_no_auth(&env, &contract_id, &Symbol::new(&env, PAUSER_ROLE), &contract_id);
+        grant_role_no_auth(&env, &contract_id, &Symbol::new(&env, UNPAUSER_ROLE), &contract_id);
+    }
+
     pub fn assert_not_paused(env: Env) {
         <Self as OFTPausableInternal>::__assert_not_paused(&env)
     }
@@ -35,6 +47,7 @@ impl PausableTestContract {
 
 struct TestSetup {
     client: PausableTestContractClient<'static>,
+    contract_id: Address,
 }
 
 fn setup() -> TestSetup {
@@ -43,7 +56,8 @@ fn setup() -> TestSetup {
 
     let contract_id = env.register(PausableTestContract, ());
     let client = PausableTestContractClient::new(&env, &contract_id);
-    TestSetup { client }
+    client.init_roles();
+    TestSetup { client, contract_id }
 }
 
 // ============================================================================
@@ -58,45 +72,45 @@ fn test_initial_state_is_not_paused() {
 }
 
 // ============================================================================
-// Set Paused Tests
+// Pause/Unpause Tests
 // ============================================================================
 
 #[test]
-fn test_set_paused_to_true() {
-    let TestSetup { client, .. } = setup();
+fn test_pause_sets_paused() {
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_paused(&true);
+    client.pause(&contract_id);
     assert!(client.is_paused());
 }
 
 #[test]
-fn test_set_paused_to_false_after_pausing() {
-    let TestSetup { client, .. } = setup();
+fn test_unpause_after_pausing() {
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_paused(&true);
+    client.pause(&contract_id);
     assert!(client.is_paused());
 
-    client.set_paused(&false);
+    client.unpause(&contract_id);
     assert!(!client.is_paused());
 }
 
 #[test]
-fn test_set_paused_rejects_same_value_when_unpaused() {
-    let TestSetup { client, .. } = setup();
+fn test_unpause_unchanged_when_unpaused() {
+    let TestSetup { client, contract_id, .. } = setup();
 
-    // Initially not paused, setting to false should fail
-    let res = client.try_set_paused(&false);
+    // Initially not paused, unpause should fail
+    let res = client.try_unpause(&contract_id);
     assert_eq!(res.err().unwrap().ok().unwrap(), OFTPausableError::PauseStatusUnchanged.into());
 }
 
 #[test]
-fn test_set_paused_rejects_same_value_when_paused() {
-    let TestSetup { client, .. } = setup();
+fn test_pause_unchanged_when_paused() {
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_paused(&true);
+    client.pause(&contract_id);
 
-    // Already paused, setting to true should fail
-    let res = client.try_set_paused(&true);
+    // Already paused, pause again should fail
+    let res = client.try_pause(&contract_id);
     assert_eq!(res.err().unwrap().ok().unwrap(), OFTPausableError::PauseStatusUnchanged.into());
 }
 
@@ -114,9 +128,9 @@ fn test_assert_not_paused_succeeds_when_not_paused() {
 
 #[test]
 fn test_assert_not_paused_fails_when_paused() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_paused(&true);
+    client.pause(&contract_id);
 
     let res = client.try_assert_not_paused();
     assert_eq!(res.err().unwrap().ok().unwrap(), OFTPausableError::Paused.into());
@@ -128,14 +142,14 @@ fn test_assert_not_paused_fails_when_paused() {
 
 #[test]
 fn test_pause_unpause_toggle_multiple_times() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
     // Toggle multiple times
     for _ in 0..3 {
         assert!(!client.is_paused());
-        client.set_paused(&true);
+        client.pause(&contract_id);
         assert!(client.is_paused());
-        client.set_paused(&false);
+        client.unpause(&contract_id);
     }
     assert!(!client.is_paused());
 }

package/contracts/oapps/oft/src/tests/extensions/rate_limiter.rs

@@ -4,8 +4,10 @@ use crate as oft;
 use crate::extensions::rate_limiter::{
     Direction, Mode, RateLimitConfig, RateLimitError, RateLimiter, RateLimiterInternal,
 };
-use soroban_sdk::{contract, contractimpl, testutils::Ledger as _, Address, Env};
+use crate::extensions::rate_limiter::RATE_LIMITER_ADMIN_ROLE;
+use soroban_sdk::{contract, contractimpl, testutils::Ledger as _, Address, Env, Symbol};
 use utils::auth::Auth;
+use utils::rbac::{grant_role_no_auth, RoleBasedAccessControl};
 
 // ============================================================================
 // Test Contract
@@ -25,8 +27,22 @@ impl RateLimiterInternal for RateLimiterTestContract {}
 #[contractimpl(contracttrait)]
 impl RateLimiter for RateLimiterTestContract {}
 
+#[contractimpl(contracttrait)]
+impl RoleBasedAccessControl for RateLimiterTestContract {}
+
 #[contractimpl]
 impl RateLimiterTestContract {
+    /// Test-only: grants RATE_LIMITER_ADMIN_ROLE to the contract.
+    pub fn init_roles(env: Env) {
+        let contract_id = env.current_contract_address();
+        grant_role_no_auth(
+            &env,
+            &contract_id,
+            &Symbol::new(&env, RATE_LIMITER_ADMIN_ROLE),
+            &contract_id,
+        );
+    }
+
     pub fn consume(env: Env, direction: Direction, eid: u32, amount: i128) {
         <Self as RateLimiterInternal>::__consume_rate_limit_capacity(&env, &direction, eid, amount)
     }
@@ -43,6 +59,7 @@ impl RateLimiterTestContract {
 struct TestSetup {
     env: Env,
     client: RateLimiterTestContractClient<'static>,
+    contract_id: Address,
 }
 
 fn setup() -> TestSetup {
@@ -53,7 +70,8 @@ fn setup() -> TestSetup {
 
     let contract_id = env.register(RateLimiterTestContract, ());
     let client = RateLimiterTestContractClient::new(&env, &contract_id);
-    TestSetup { env, client }
+    client.init_roles();
+    TestSetup { env, client, contract_id }
 }
 
 fn config(limit: i128, window_seconds: u64) -> RateLimitConfig {
@@ -93,29 +111,29 @@ fn test_no_config_consume_and_release_are_noop() {
 
 #[test]
 fn test_set_rate_limit_rejects_invalid_config() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
     // Invalid: window_seconds == 0
-    let res = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(config(10, 0)));
+    let res = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(config(10, 0)), &contract_id);
     assert_eq!(res.err().unwrap().ok().unwrap(), RateLimitError::InvalidConfig.into());
 
     // Invalid: limit < 0
-    let res2 = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(config(-1, 10)));
+    let res2 = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(config(-1, 10)), &contract_id);
     assert_eq!(res2.err().unwrap().ok().unwrap(), RateLimitError::InvalidConfig.into());
 }
 
 #[test]
 fn test_set_rate_limit_rejects_same_value() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
     // SameValue when removing a non-existent config.
-    let res = client.try_set_rate_limit(&Direction::Outbound, &1u32, &None);
+    let res = client.try_set_rate_limit(&Direction::Outbound, &1u32, &None, &contract_id);
     assert_eq!(res.err().unwrap().ok().unwrap(), RateLimitError::SameValue.into());
 
     // Set a valid config; setting same again should be SameValue.
     let cfg = config(100, 10);
-    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(cfg.clone()));
-    let res2 = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(cfg));
+    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(cfg.clone()), &contract_id);
+    let res2 = client.try_set_rate_limit(&Direction::Outbound, &1u32, &Some(cfg), &contract_id);
     assert_eq!(res2.err().unwrap().ok().unwrap(), RateLimitError::SameValue.into());
 }
 
@@ -125,9 +143,9 @@ fn test_set_rate_limit_rejects_same_value() {
 
 #[test]
 fn test_consume_rejects_negative_amount() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(config(100, 100)));
+    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(config(100, 100)), &contract_id);
 
     let res = client.try_consume(&Direction::Outbound, &1u32, &-1i128);
     assert_eq!(res.err().unwrap().ok().unwrap(), RateLimitError::InvalidAmount.into());
@@ -135,9 +153,9 @@ fn test_consume_rejects_negative_amount() {
 
 #[test]
 fn test_release_rejects_negative_amount() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
 
-    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(config(100, 100)));
+    client.set_rate_limit(&Direction::Outbound, &1u32, &Some(config(100, 100)), &contract_id);
 
     let res = client.try_release(&Direction::Outbound, &1u32, &-1i128);
     assert_eq!(res.err().unwrap().ok().unwrap(), RateLimitError::InvalidAmount.into());
@@ -145,11 +163,11 @@ fn test_release_rejects_negative_amount() {
 
 #[test]
 fn test_consume_within_capacity() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
 
     client.consume(&direction, &eid, &60i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60);
@@ -158,11 +176,11 @@ fn test_consume_within_capacity() {
 
 #[test]
 fn test_consume_exceeds_capacity() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
     client.consume(&direction, &eid, &60i128);
 
     let res = client.try_consume(&direction, &eid, &41i128);
@@ -171,11 +189,11 @@ fn test_consume_exceeds_capacity() {
 
 #[test]
 fn test_release_reduces_in_flight() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
     client.consume(&direction, &eid, &60i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60);
 
@@ -194,12 +212,12 @@ fn test_release_reduces_in_flight() {
 
 #[test]
 fn test_decay_reduces_in_flight_over_time() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 5u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
 
     client.consume(&direction, &eid, &100i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 100);
@@ -212,12 +230,12 @@ fn test_decay_reduces_in_flight_over_time() {
 
 #[test]
 fn test_timestamp_regression_errors() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 5u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
     client.consume(&direction, &eid, &100i128);
 
     // Set timestamp backwards => should error.
@@ -232,17 +250,17 @@ fn test_timestamp_regression_errors() {
 
 #[test]
 fn test_updating_config_checkpoints_in_flight() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 42u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
     client.consume(&direction, &eid, &100i128);
 
     // Move forward, then update config.
     env.ledger().set_timestamp(1_005);
-    client.set_rate_limit(&direction, &eid, &Some(config(200, 20)));
+    client.set_rate_limit(&direction, &eid, &Some(config(200, 20)), &contract_id);
 
     // At t=1005, old config would have decayed to 50.
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 50);
@@ -255,29 +273,29 @@ fn test_updating_config_checkpoints_in_flight() {
 
 #[test]
 fn test_reducing_limit_below_in_flight_clamps_capacity() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Inbound;
     let eid = 123u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
     client.consume(&direction, &eid, &90i128);
 
     // Reduce limit below in-flight.
-    client.set_rate_limit(&direction, &eid, &Some(config(50, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(50, 100)), &contract_id);
     assert_eq!(client.rate_limit_capacity(&direction, &eid), 0);
 }
 
 #[test]
 fn test_removing_rate_limit_makes_unlimited() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Inbound;
     let eid = 321u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(10, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(10, 10)), &contract_id);
     assert_eq!(client.rate_limit_capacity(&direction, &eid), 10);
 
-    client.set_rate_limit(&direction, &eid, &None);
+    client.set_rate_limit(&direction, &eid, &None, &contract_id);
     assert_eq!(client.rate_limit_capacity(&direction, &eid), i128::MAX);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 0);
 }
@@ -296,12 +314,12 @@ fn test_rate_limit_config_returns_none_when_unset() {
 
 #[test]
 fn test_rate_limit_config_returns_config_when_set() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 42u32;
 
     let cfg = config(500, 60);
-    client.set_rate_limit(&direction, &eid, &Some(cfg.clone()));
+    client.set_rate_limit(&direction, &eid, &Some(cfg.clone()), &contract_id);
 
     let result = client.rate_limit_config(&direction, &eid);
     assert_eq!(result, Some(cfg));
@@ -309,14 +327,14 @@ fn test_rate_limit_config_returns_config_when_set() {
 
 #[test]
 fn test_rate_limit_config_returns_none_after_removal() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Inbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
     assert!(client.rate_limit_config(&direction, &eid).is_some());
 
-    client.set_rate_limit(&direction, &eid, &None);
+    client.set_rate_limit(&direction, &eid, &None, &contract_id);
     assert_eq!(client.rate_limit_config(&direction, &eid), None);
 }
 
@@ -326,12 +344,12 @@ fn test_rate_limit_config_returns_none_after_removal() {
 
 #[test]
 fn test_inbound_outbound_are_independent() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let eid = 10u32;
 
     // Set different configs for inbound and outbound
-    client.set_rate_limit(&Direction::Inbound, &eid, &Some(config(100, 10)));
-    client.set_rate_limit(&Direction::Outbound, &eid, &Some(config(200, 20)));
+    client.set_rate_limit(&Direction::Inbound, &eid, &Some(config(100, 10)), &contract_id);
+    client.set_rate_limit(&Direction::Outbound, &eid, &Some(config(200, 20)), &contract_id);
 
     assert_eq!(client.rate_limit_config(&Direction::Inbound, &eid), Some(config(100, 10)));
     assert_eq!(client.rate_limit_config(&Direction::Outbound, &eid), Some(config(200, 20)));
@@ -344,11 +362,11 @@ fn test_inbound_outbound_are_independent() {
 
 #[test]
 fn test_different_eids_are_independent() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
 
-    client.set_rate_limit(&direction, &1u32, &Some(config(100, 10)));
-    client.set_rate_limit(&direction, &2u32, &Some(config(200, 20)));
+    client.set_rate_limit(&direction, &1u32, &Some(config(100, 10)), &contract_id);
+    client.set_rate_limit(&direction, &2u32, &Some(config(200, 20)), &contract_id);
 
     client.consume(&direction, &1u32, &50i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &1u32), 50);
@@ -361,11 +379,11 @@ fn test_different_eids_are_independent() {
 
 #[test]
 fn test_gross_mode_does_not_release_capacity() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)), &contract_id);
     client.consume(&direction, &eid, &60i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60);
 
@@ -376,11 +394,11 @@ fn test_gross_mode_does_not_release_capacity() {
 
 #[test]
 fn test_net_mode_does_release_capacity() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 77u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)), &contract_id);
     client.consume(&direction, &eid, &60i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60);
 
@@ -391,11 +409,11 @@ fn test_net_mode_does_release_capacity() {
 
 #[test]
 fn test_gross_mode_enforces_absolute_limit() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)), &contract_id);
 
     // Consume 80
     client.consume(&direction, &eid, &80i128);
@@ -418,12 +436,12 @@ fn test_gross_mode_enforces_absolute_limit() {
 
 #[test]
 fn test_gross_mode_decay_still_works() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 10, Mode::Gross)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 10, Mode::Gross)), &contract_id);
 
     // Consume full capacity
     client.consume(&direction, &eid, &100i128);
@@ -441,18 +459,18 @@ fn test_gross_mode_decay_still_works() {
 
 #[test]
 fn test_switching_from_net_to_gross_mode() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     // Start with Net mode
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)), &contract_id);
     client.consume(&direction, &eid, &60i128);
     client.release(&direction, &eid, &20i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 40, "Net mode should release");
 
     // Switch to Gross mode
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)), &contract_id);
 
     // in-flight should be checkpointed
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 40);
@@ -464,18 +482,18 @@ fn test_switching_from_net_to_gross_mode() {
 
 #[test]
 fn test_switching_from_gross_to_net_mode() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     // Start with Gross mode
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Gross)), &contract_id);
     client.consume(&direction, &eid, &60i128);
     client.release(&direction, &eid, &20i128);
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60, "Gross mode should not release");
 
     // Switch to Net mode
-    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)));
+    client.set_rate_limit(&direction, &eid, &Some(config_with_mode(100, 100, Mode::Net)), &contract_id);
 
     // in-flight should be checkpointed
     assert_eq!(client.rate_limit_in_flight(&direction, &eid), 60);
@@ -491,11 +509,11 @@ fn test_switching_from_gross_to_net_mode() {
 
 #[test]
 fn test_consume_zero_amount_succeeds() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
 
     // Consuming zero should succeed
     client.consume(&direction, &eid, &0i128);
@@ -504,11 +522,11 @@ fn test_consume_zero_amount_succeeds() {
 
 #[test]
 fn test_release_zero_amount_succeeds() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
     client.consume(&direction, &eid, &50i128);
 
     // Releasing zero should succeed
@@ -518,11 +536,11 @@ fn test_release_zero_amount_succeeds() {
 
 #[test]
 fn test_consume_exact_capacity() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 100)), &contract_id);
 
     // Consuming exact capacity should succeed
     client.consume(&direction, &eid, &100i128);
@@ -532,12 +550,12 @@ fn test_consume_exact_capacity() {
 
 #[test]
 fn test_decay_fully_restores_capacity() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
 
     client.consume(&direction, &eid, &100i128);
     assert_eq!(client.rate_limit_capacity(&direction, &eid), 0);
@@ -550,12 +568,12 @@ fn test_decay_fully_restores_capacity() {
 
 #[test]
 fn test_decay_beyond_window_clamps_to_zero() {
-    let TestSetup { env, client } = setup();
+    let TestSetup { env, client, contract_id } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     env.ledger().set_timestamp(1_000);
-    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(100, 10)), &contract_id);
 
     client.consume(&direction, &eid, &50i128);
 
@@ -566,12 +584,12 @@ fn test_decay_beyond_window_clamps_to_zero() {
 
 #[test]
 fn test_limit_zero_blocks_all_consumption() {
-    let TestSetup { client, .. } = setup();
+    let TestSetup { client, contract_id, .. } = setup();
     let direction = Direction::Outbound;
     let eid = 1u32;
 
     // Set limit to 0 (valid config)
-    client.set_rate_limit(&direction, &eid, &Some(config(0, 10)));
+    client.set_rate_limit(&direction, &eid, &Some(config(0, 10)), &contract_id);
     assert_eq!(client.rate_limit_capacity(&direction, &eid), 0);
 
     // Any non-zero consumption should fail