feat(payment-escrow): Implement treasury with admin withdrawal

contracts/payment_escrow/src/errors.rs

@@ -31,4 +31,6 @@ pub enum Error {
     PaymentTokenNotSet = 12,
     /// Fee recipient address has not been set.
     FeeRecipientNotSet = 13,
+    InsufficientBalance = 14,
+}
 }

contracts/payment_escrow/src/lib.rs

@@ -418,12 +418,12 @@ impl PaymentEscrowContract {
         let token_client = token::Client::new(&env, &escrow.payment_token);
 
         if release_to_beneficiary {
-            // Transfer the fee, if any
+            // Add fee to treasury
             if escrow.fee_amount > 0 {
-                token_client.transfer(
-                    &env.current_contract_address(),
-                    &escrow.fee_recipient,
-                    &escrow.fee_amount,
+                let treasury = Self::get_treasury_amount(&env);
+                env.storage().instance().set(
+                    &DataKey::TreasuryAmount,
+                    &(treasury + escrow.fee_amount),
                 );
             }
 
@@ -562,4 +562,42 @@ impl PaymentEscrowContract {
     pub fn dispute_window(env: Env) -> u64 {
         Self::get_dispute_window(&env)
     }
+
+    // ── Treasury ──────────────────────────────────────────────────────────────
+
+    /// Withdraw accumulated fees from the treasury.
+    pub fn withdraw_treasury(
+        env: Env,
+        caller: Address,
+        recipient: Address,
+        amount: i128,
+    ) -> Result<(), Error> {
+        Self::require_admin(&env, &caller)?;
+
+        let treasury = Self::get_treasury_amount(&env);
+        if amount > treasury {
+            return Err(Error::InsufficientBalance);
+        }
+
+        let payment_token = Self::get_payment_token(&env)?;
+        let token_client = token::Client::new(&env, &payment_token);
+
+        token_client.transfer(&env.current_contract_address(), &recipient, &amount);
+
+        env.storage()
+            .instance()
+            .set(&DataKey::TreasuryAmount, &(treasury - amount));
+
+        env.events().publish(
+            (symbol_short!("treasury_w"),),
+            (recipient, amount, env.ledger().timestamp()),
+        );
+        Ok(())
+    }
+
+    /// Return the current treasury balance.
+    pub fn treasury_balance(env: Env) -> i128 {
+        Self::get_treasury_amount(&env)
+    }
+}
 }

contracts/payment_escrow/src/treasury.rs

@@ -0,0 +1,154 @@
+"""// contracts/payment_escrow/src/treasury.rs
+#![cfg(test)]
+
+use soroban_sdk::{
+    testutils::{Address as _, Events},
+    Address, Env, String,
+};
+
+use crate::{
+    test::{
+        helpers::{
+            create_and_initialize_contract, create_escrow, create_token, get_ledger_timestamp,
+        },
+        setup::Setup,
+    },
+    Error, PaymentEscrowContract,
+};
+
+#[test]
+fn test_withdraw_treasury_unauthorized() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let setup = Setup::new(&env);
+    let depositor = Address::generate(&env);
+    let unauthorized_caller = Address::generate(&env);
+
+    create_and_initialize_contract(
+        &env,
+        &setup.contract_id,
+        &setup.admin,
+        &setup.token.address,
+        10,
+        &setup.fee_recipient,
+        100,
+    );
+
+    let escrow_id = String::from_str(&env, "escrow-1");
+    create_escrow(
+        &env,
+        &setup.contract_id,
+        &setup.token.address,
+        &depositor,
+        escrow_id.clone(),
+        &setup.beneficiary,
+        1000,
+        "Test Escrow",
+        0,
+    );
+
+    let res = PaymentEscrowContract::new(&env, &setup.contract_id).try_withdraw_treasury(
+        &unauthorized_caller,
+        &setup.fee_recipient,
+        10,
+    );
+    assert_eq!(res, Err(Ok(Error::Unauthorized)));
+}
+
+#[test]
+fn test_withdraw_treasury_insufficient_balance() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let setup = Setup::new(&env);
+    let depositor = Address::generate(&env);
+
+    create_and_initialize_contract(
+        &env,
+        &setup.contract_id,
+        &setup.admin,
+        &setup.token.address,
+        10,
+        &setup.fee_recipient,
+        100,
+    );
+
+    let escrow_id = String::from_str(&env, "escrow-1");
+    create_escrow(
+        &env,
+        &setup.contract_id,
+        &setup.token.address,
+        &depositor,
+        escrow_id.clone(),
+        &setup.beneficiary,
+        1000,
+        "Test Escrow",
+        0,
+    );
+
+    let res = PaymentEscrowContract::new(&env, &setup.contract_id).try_withdraw_treasury(
+        &setup.admin,
+        &setup.fee_recipient,
+        10,
+    );
+    assert_eq!(res, Err(Ok(Error::InsufficientBalance)));
+}
+
+#[test]
+fn test_withdraw_treasury_success() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let setup = Setup::new(&env);
+    let depositor = Address::generate(&env);
+    let contract = PaymentEscrowContract::new(&env, &setup.contract_id);
+
+    create_and_initialize_contract(
+        &env,
+        &setup.contract_id,
+        &setup.admin,
+        &setup.token.address,
+        10,
+        &setup.fee_recipient,
+        100, // 1% fee
+    );
+
+    // Create and release an escrow to generate fees
+    let escrow_id = String::from_str(&env, "escrow-1");
+    create_escrow(
+        &env,
+        &setup.contract_id,
+        &setup.token.address,
+        &depositor,
+        escrow_id.clone(),
+        &setup.beneficiary,
+        1000,
+        "Test Escrow",
+        0,
+    );
+    contract.release(&setup.admin, escrow_id);
+
+    // Withdraw a portion of the treasury
+    contract.withdraw_treasury(&setup.admin, &setup.fee_recipient, 5);
+    assert_eq!(setup.token.balance(&setup.fee_recipient), 5);
+    assert_eq!(setup.token.balance(&setup.contract_id), 995);
+
+    // Withdraw the rest
+    contract.withdraw_treasury(&setup.admin, &setup.fee_recipient, 5);
+    assert_eq!(setup.token.balance(&setup.fee_recipient), 10);
+    assert_eq!(setup.token.balance(&setup.contract_id), 990);
+
+    // Check events
+    let event = env.events().all().last().unwrap();
+    let timestamp = get_ledger_timestamp(&env);
+    assert_eq!(
+        event,
+        (
+            setup.contract_id.clone(),
+            ("treasury_w",),
+            (setup.fee_recipient.clone(), 5i128, timestamp).into_val(&env)
+        )
+    );
+}
+""