Merge pull request #1544 from martinshub-tech/fix-admin-force-rebalance-tests

Fix admin force rebalance tests

Author: ritik4ever

contracts/benchmarks/BASELINES.md

@@ -7,6 +7,7 @@ Current baseline (native test runtime):
 - `initialize`: CPU `1,500,000`, memory `200,000`
 - `create_portfolio`: CPU `2,500,000`, memory `300,000`
 - `execute_rebalance`: CPU `5,000,000`, memory `500,000`
+- `execute_rebalance_max_assets`: CPU `10,000,000`, memory `1,000,000`
 - `deposit`: CPU `2,000,000`, memory `250,000`
 
 The benchmark assertions fail if any metric exceeds 120% of the baseline.

contracts/src/test.rs

@@ -44,6 +44,8 @@ const BASELINE_CREATE_PORTFOLIO_CPU: u64 = 2_500_000;
 const BASELINE_CREATE_PORTFOLIO_MEM: u64 = 300_000;
 const BASELINE_EXECUTE_REBALANCE_CPU: u64 = 5_000_000;
 const BASELINE_EXECUTE_REBALANCE_MEM: u64 = 500_000;
+const BASELINE_EXECUTE_REBALANCE_MAX_ASSETS_CPU: u64 = 10_000_000; // Will adjust later
+const BASELINE_EXECUTE_REBALANCE_MAX_ASSETS_MEM: u64 = 1_000_000; // Will adjust later
 const BASELINE_DEPOSIT_CPU: u64 = 2_000_000;
 const BASELINE_DEPOSIT_MEM: u64 = 250_000;
 
@@ -2089,6 +2091,68 @@ fn test_admin_force_rebalance_bypasses_cooldown() {
     assert_eq!(portfolio.last_rebalance, 10010);
 }
 
+#[test]
+#[should_panic]
+fn test_admin_force_rebalance_non_admin_rejected() {
+    let env = Env::default();
+    env.mock_all_auths();
+    let contract_id = env.register_contract(None, PortfolioRebalancer);
+    let client = PortfolioRebalancerClient::new(&env, &contract_id);
+    let reflector_id = env.register_contract(None, reflector_contract::MockReflector);
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    client.initialize(&admin, &reflector_id);
+
+    let mut allocations = Map::new(&env);
+    let asset = Address::generate(&env);
+    allocations.set(asset, 10000);
+    let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
+
+    let unauthorized = Address::generate(&env);
+    let actual_balances = Map::new(&env);
+    client
+        .mock_auths(&[MockAuth {
+            address: &unauthorized,
+            invoke: &MockAuthInvoke {
+                contract: &contract_id,
+                fn_name: "admin_force_rebalance",
+                args: (pid, actual_balances.clone()).into_val(&env),
+                sub_invokes: &[],
+            },
+        }])
+        .admin_force_rebalance(&pid, &actual_balances);
+}
+
+#[test]
+fn test_admin_force_rebalance_admin_success() {
+    let env = Env::default();
+    env.mock_all_auths();
+    let contract_id = env.register_contract(None, PortfolioRebalancer);
+    let client = PortfolioRebalancerClient::new(&env, &contract_id);
+    let reflector_id = env.register_contract(None, reflector_contract::MockReflector);
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    client.initialize(&admin, &reflector_id);
+
+    let mut allocations = Map::new(&env);
+    let asset = Address::generate(&env);
+    allocations.set(asset, 10000);
+    let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
+
+    let actual_balances = Map::new(&env);
+    client
+        .mock_auths(&[MockAuth {
+            address: &admin,
+            invoke: &MockAuthInvoke {
+                contract: &contract_id,
+                fn_name: "admin_force_rebalance",
+                args: (pid, actual_balances.clone()).into_val(&env),
+                sub_invokes: &[],
+            },
+        }])
+        .admin_force_rebalance(&pid, &actual_balances);
+}
+
 #[test]
 fn test_portfolio_invariants_helper_rejects_invalid_allocations() {
     let env = Env::default();
@@ -2262,6 +2326,74 @@ fn benchmark_deposit_gas() {
     );
 }
 
+#[test]
+fn benchmark_execute_rebalance_max_assets() {
+    let env = Env::default();
+    env.mock_all_auths();
+    env.budget().reset_unlimited();
+
+    let contract_id = env.register_contract(None, PortfolioRebalancer);
+    let client = PortfolioRebalancerClient::new(&env, &contract_id);
+    let reflector_id = env.register_contract(None, reflector_contract::MockReflector);
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    let _ = client.initialize(&admin, &reflector_id);
+
+    let mut allocations = Map::new(&env);
+    for _ in 0..MAX_PORTFOLIO_ASSETS {
+        allocations.set(Address::generate(&env), ALLOCATION_DENOMINATOR / MAX_PORTFOLIO_ASSETS);
+    }
+    let pid = create_portfolio_with_defaults(&env, &client, &user, &allocations, 5, 50);
+
+    env.ledger().with_mut(|li| {
+        li.timestamp = 15_000;
+    });
+
+    env.budget().reset_tracker();
+    let _ = client.execute_rebalance(&pid, &Map::new(&env));
+    
+    std::println!("CPU used for max assets: {}", env.budget().cpu_instruction_cost());
+    std::println!("MEM used for max assets: {}", env.budget().memory_bytes_cost());
+    
+    assert_cost_within_tolerance(
+        "execute_rebalance_max_assets",
+        env.budget().cpu_instruction_cost(),
+        env.budget().memory_bytes_cost(),
+        BASELINE_EXECUTE_REBALANCE_MAX_ASSETS_CPU,
+        BASELINE_EXECUTE_REBALANCE_MAX_ASSETS_MEM,
+    );
+}
+
+#[test]
+fn test_execute_rebalance_max_assets_plus_one_rejected() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let contract_id = env.register_contract(None, PortfolioRebalancer);
+    let client = PortfolioRebalancerClient::new(&env, &contract_id);
+    let reflector_id = env.register_contract(None, reflector_contract::MockReflector);
+    let admin = Address::generate(&env);
+    let user = Address::generate(&env);
+    let _ = client.initialize(&admin, &reflector_id);
+
+    let mut too_many_allocations = Map::new(&env);
+    for _ in 0..(MAX_PORTFOLIO_ASSETS + 1) {
+        too_many_allocations.set(Address::generate(&env), ALLOCATION_DENOMINATOR / (MAX_PORTFOLIO_ASSETS + 1));
+    }
+    
+    let too_many_decimals = allocation_decimals(&env, &too_many_allocations, DEFAULT_ASSET_DECIMALS);
+    
+    let result = client.try_create_portfolio(
+        &user,
+        &too_many_allocations,
+        &too_many_decimals,
+        &5,
+        &50,
+        &CURRENT_SLIPPAGE_POLICY_VERSION,
+    );
+    assert_eq!(result, Err(Ok(Error::TooManyAssets)));
+}
+
 // ── Issue #861: rebalance validates allocation sum ──────────────────────
 
 #[test]