@layerzerolabs/protocol-starknet-v2 0.0.34

package/layerzero/tests/oapps/test_oapp_core.cairo

@@ -0,0 +1,750 @@
+//! OApp core tests
+
+use core::num::traits::Pow;
+use layerzero::common::structs::packet::Origin;
+use layerzero::endpoint::interfaces::layerzero_receiver::{
+    ILayerZeroReceiverDispatcher, ILayerZeroReceiverDispatcherTrait,
+    ILayerZeroReceiverSafeDispatcher, ILayerZeroReceiverSafeDispatcherTrait,
+};
+use layerzero::oapps::oapp::errors::{
+    err_lz_token_unavailable, err_no_peer, err_not_enough_lz_token,
+    err_not_enough_lz_token_allowance, err_not_enough_native, err_not_enough_native_allowance,
+    err_only_endpoint, err_only_peer,
+};
+use layerzero::oapps::oapp::events::PeerSet;
+use layerzero::oapps::oapp::interface::{
+    IOAppDispatcher, IOAppDispatcherTrait, IOAppReceiverDispatcher, IOAppReceiverDispatcherTrait,
+    IOAppSafeDispatcher, IOAppSafeDispatcherTrait,
+};
+use layerzero::oapps::oapp::oapp_core::OAppCoreComponent;
+use layerzero::{MessageReceipt, MessagingFee};
+use lz_utils::bytes::Bytes32;
+use openzeppelin::access::ownable::OwnableComponent;
+use openzeppelin::access::ownable::interface::{
+    IOwnableDispatcher, IOwnableDispatcherTrait, IOwnableSafeDispatcher,
+};
+use openzeppelin::token::erc20::interface::{IERC20Dispatcher, IERC20DispatcherTrait};
+use snforge_std::{
+    ContractClassTrait, DeclareResultTrait, EventSpyAssertionsTrait, declare, spy_events,
+    start_cheat_caller_address, start_mock_call, stop_cheat_caller_address,
+};
+use starknet::ContractAddress;
+use starkware_utils_testing::test_utils::{assert_panic_with_error, assert_panic_with_felt_error};
+use crate::mocks::oapp_core::interface::{
+    IMockOAppCoreDispatcher, IMockOAppCoreDispatcherTrait, IMockOAppCoreSafeDispatcher,
+    IMockOAppCoreSafeDispatcherTrait,
+};
+use crate::mocks::oapp_core::oapp_core::MockOAppCore;
+
+// Test constants
+const FAKE_ENDPOINT: ContractAddress = 'fake_endpoint'.try_into().unwrap();
+const OWNER: ContractAddress = 'owner'.try_into().unwrap();
+const NON_OWNER: ContractAddress = 'non_owner'.try_into().unwrap();
+const EXECUTOR: ContractAddress = 'executor'.try_into().unwrap();
+const ZERO_ADDRESS: ContractAddress = 0.try_into().unwrap();
+const EID_1: u32 = 101;
+const EID_2: u32 = 102;
+const SRC_EID: u32 = 201;
+const UNKNOWN_EID: u32 = 999;
+const INITIAL_SUPPLY: u256 = 10_u256.pow(10);
+
+fn MOCK_PEER_1() -> Bytes32 {
+    Bytes32 { value: 0x123456789abcdef }
+}
+
+fn MOCK_PEER_2() -> Bytes32 {
+    Bytes32 { value: 0xfedcba987654321 }
+}
+
+fn ZERO_PEER() -> Bytes32 {
+    Bytes32 { value: 0 }
+}
+
+fn MOCK_GUID() -> Bytes32 {
+    Bytes32 { value: 0xabcdef123456789 }
+}
+
+fn MOCK_MESSAGE() -> ByteArray {
+    "test message"
+}
+
+fn MOCK_ORIGIN() -> Origin {
+    Origin { src_eid: SRC_EID, sender: MOCK_PEER_1(), nonce: 1 }
+}
+
+#[derive(Drop, Copy)]
+struct TestSetup {
+    oapp_address: ContractAddress,
+    oapp_dispatcher: IOAppDispatcher,
+    ownable_dispatcher: IOwnableDispatcher,
+    lz_receiver_dispatcher: ILayerZeroReceiverDispatcher,
+    oapp_receiver_dispatcher: IOAppReceiverDispatcher,
+    oapp_safe_dispatcher: IOAppSafeDispatcher,
+    ownable_safe_dispatcher: IOwnableSafeDispatcher,
+    lz_receiver_safe_dispatcher: ILayerZeroReceiverSafeDispatcher,
+    mock_oapp_core_dispatcher: IMockOAppCoreDispatcher,
+    mock_oapp_core_safe_dispatcher: IMockOAppCoreSafeDispatcher,
+    stark_token_dispatcher: IERC20Dispatcher,
+    lz_token_dispatcher: IERC20Dispatcher,
+}
+
+
+fn deploy_test_erc20(owner: ContractAddress) -> ContractAddress {
+    let contract = declare("MockERC20").unwrap().contract_class();
+    let mut params = array![];
+    INITIAL_SUPPLY.serialize(ref params);
+    owner.serialize(ref params);
+
+    let (address, _) = contract.deploy(@params).unwrap();
+    address
+}
+
+fn deploy_test_oapp_core(
+    endpoint: ContractAddress, owner: ContractAddress, stark_token: ContractAddress,
+) -> ContractAddress {
+    let contract = declare("MockOAppCore").unwrap().contract_class();
+
+    let (address, _) = contract
+        .deploy(@array![endpoint.into(), owner.into(), stark_token.into()])
+        .unwrap();
+    address
+}
+
+fn setup() -> TestSetup {
+    let stark_token = deploy_test_erc20(OWNER);
+    let lz_token = deploy_test_erc20(OWNER);
+    let oapp_address = deploy_test_oapp_core(FAKE_ENDPOINT, OWNER, stark_token);
+    start_mock_call(FAKE_ENDPOINT, selector!("get_lz_token"), lz_token);
+
+    let oapp_dispatcher = IOAppDispatcher { contract_address: oapp_address };
+    let ownable_dispatcher = IOwnableDispatcher { contract_address: oapp_address };
+    let lz_receiver_dispatcher = ILayerZeroReceiverDispatcher { contract_address: oapp_address };
+    let oapp_receiver_dispatcher = IOAppReceiverDispatcher { contract_address: oapp_address };
+    let oapp_safe_dispatcher = IOAppSafeDispatcher { contract_address: oapp_address };
+    let ownable_safe_dispatcher = IOwnableSafeDispatcher { contract_address: oapp_address };
+    let lz_receiver_safe_dispatcher = ILayerZeroReceiverSafeDispatcher {
+        contract_address: oapp_address,
+    };
+    let mock_oapp_core_dispatcher = IMockOAppCoreDispatcher { contract_address: oapp_address };
+    let mock_oapp_core_safe_dispatcher = IMockOAppCoreSafeDispatcher {
+        contract_address: oapp_address,
+    };
+    let stark_token_dispatcher = IERC20Dispatcher { contract_address: stark_token };
+    let lz_token_dispatcher = IERC20Dispatcher { contract_address: lz_token };
+    // Set up peer for testing
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(SRC_EID, MOCK_PEER_1());
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    TestSetup {
+        oapp_address,
+        oapp_dispatcher,
+        ownable_dispatcher,
+        lz_receiver_dispatcher,
+        oapp_receiver_dispatcher,
+        oapp_safe_dispatcher,
+        ownable_safe_dispatcher,
+        lz_receiver_safe_dispatcher,
+        mock_oapp_core_dispatcher,
+        mock_oapp_core_safe_dispatcher,
+        stark_token_dispatcher,
+        lz_token_dispatcher,
+    }
+}
+
+// =============================== Test Constructor =================================
+
+#[test]
+fn test_constructor() {
+    let TestSetup { oapp_dispatcher, ownable_dispatcher, .. } = setup();
+    assert(oapp_dispatcher.get_endpoint() == FAKE_ENDPOINT, 'Endpoint not set correctly');
+    assert(ownable_dispatcher.owner() == OWNER, 'Owner not set correctly');
+}
+
+// =============================== Test OApp Version =================================
+
+#[test]
+fn test_oapp_version() {
+    let TestSetup { oapp_dispatcher, .. } = setup();
+    let (sender_version, receiver_version) = oapp_dispatcher.oapp_version();
+    assert(sender_version == 1, 'Sender version should be 1');
+    assert(receiver_version == 1, 'Receiver version should be 2');
+}
+
+// =============================== Test Set Peer =================================
+
+#[test]
+fn test_set_peer_as_owner() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+    let mut spy = spy_events();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'Peer not set correctly');
+
+    let expected_event = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_1, peer: MOCK_PEER_1() }),
+    );
+    spy.assert_emitted(@array![(oapp_address, expected_event)]);
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_set_peer_fails_as_non_owner() {
+    let TestSetup { oapp_address, oapp_safe_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, NON_OWNER);
+    let result = oapp_safe_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_felt_error(result, OwnableComponent::Errors::NOT_OWNER);
+}
+
+// =============================== Test Get Peer =================================
+
+#[test]
+fn test_get_unset_peer() {
+    let TestSetup { oapp_dispatcher, .. } = setup();
+    assert(oapp_dispatcher.get_peer(UNKNOWN_EID) == ZERO_PEER(), 'Unset peer should be zero');
+}
+
+#[test]
+fn test_get_set_peer() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_2, MOCK_PEER_2());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_2) == MOCK_PEER_2(), 'Peer should match');
+}
+
+// =============================== Test Multiple Peers =================================
+
+#[test]
+fn test_set_multiple_peers() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+    let mut spy = spy_events();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    oapp_dispatcher.set_peer(EID_2, MOCK_PEER_2());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'Peer 1 not set correctly');
+    assert(oapp_dispatcher.get_peer(EID_2) == MOCK_PEER_2(), 'Peer 2 not set correctly');
+
+    let expected_event_1 = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_1, peer: MOCK_PEER_1() }),
+    );
+    let expected_event_2 = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_2, peer: MOCK_PEER_2() }),
+    );
+    spy.assert_emitted(@array![(oapp_address, expected_event_1), (oapp_address, expected_event_2)]);
+}
+
+// =============================== Test Update Peer =================================
+
+#[test]
+fn test_update_peer() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+    let mut spy = spy_events();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'Peer not set correctly');
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_2());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_2(), 'Peer not updated correctly');
+
+    let expected_event_1 = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_1, peer: MOCK_PEER_1() }),
+    );
+    let expected_event_2 = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_1, peer: MOCK_PEER_2() }),
+    );
+    spy.assert_emitted(@array![(oapp_address, expected_event_1), (oapp_address, expected_event_2)]);
+}
+
+// =============================== Test Remove Peer =================================
+
+#[test]
+fn test_remove_peer() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+    let mut spy = spy_events();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'Peer not set correctly');
+    oapp_dispatcher.set_peer(EID_1, ZERO_PEER());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == ZERO_PEER(), 'Peer not removed correctly');
+
+    let expected_removal_event = MockOAppCore::Event::OAppCoreEvent(
+        OAppCoreComponent::Event::PeerSet(PeerSet { eid: EID_1, peer: ZERO_PEER() }),
+    );
+    spy.assert_emitted(@array![(oapp_address, expected_removal_event)]);
+}
+
+// =============================== Test Peer Independence =================================
+
+#[test]
+fn test_peer_independence() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'EID_1 peer not set');
+    assert(oapp_dispatcher.get_peer(EID_2) == ZERO_PEER(), 'EID_2 should be unset');
+    assert(oapp_dispatcher.get_peer(UNKNOWN_EID) == ZERO_PEER(), 'Other EIDs should be unset');
+}
+
+// =============================== Test Ownership =================================
+
+#[test]
+fn test_ownership_initial_setup() {
+    let TestSetup { ownable_dispatcher, .. } = setup();
+    assert(ownable_dispatcher.owner() == OWNER, 'Initial owner incorrect');
+}
+
+#[test]
+fn test_transfer_ownership() {
+    let TestSetup { oapp_address, oapp_dispatcher, ownable_dispatcher, .. } = setup();
+    let new_owner = 'new_owner'.try_into().unwrap();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    ownable_dispatcher.transfer_ownership(new_owner);
+    stop_cheat_caller_address(oapp_address);
+
+    assert(ownable_dispatcher.owner() == new_owner, 'Ownership not transferred');
+
+    start_cheat_caller_address(oapp_address, new_owner);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    assert(oapp_dispatcher.get_peer(EID_1) == MOCK_PEER_1(), 'New owner should set peer');
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_old_owner_cannot_set_peer_after_transfer() {
+    let TestSetup { oapp_address, ownable_dispatcher, oapp_safe_dispatcher, .. } = setup();
+    let new_owner = 'new_owner'.try_into().unwrap();
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    ownable_dispatcher.transfer_ownership(new_owner);
+    stop_cheat_caller_address(oapp_address);
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    let result = oapp_safe_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_felt_error(result, OwnableComponent::Errors::NOT_OWNER);
+}
+
+// =============================== Test LayerZero Receiver =================================
+
+#[test]
+fn test_lz_receive_success() {
+    let TestSetup { oapp_address, lz_receiver_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+    lz_receiver_dispatcher
+        .lz_receive(MOCK_ORIGIN(), MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_receive_fails_non_endpoint() {
+    let TestSetup { oapp_address, lz_receiver_safe_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, NON_OWNER);
+    let result = lz_receiver_safe_dispatcher
+        .lz_receive(MOCK_ORIGIN(), MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_error(result, err_only_endpoint(FAKE_ENDPOINT));
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_receive_fails_unknown_peer() {
+    let TestSetup { oapp_address, lz_receiver_safe_dispatcher, .. } = setup();
+    let unknown_origin = Origin {
+        src_eid: UNKNOWN_EID, sender: Bytes32 { value: 0x999 }, nonce: 1,
+    };
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+    let result = lz_receiver_safe_dispatcher
+        .lz_receive(unknown_origin, MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_error(result, err_no_peer(UNKNOWN_EID));
+}
+
+#[test]
+fn test_allow_initialize_path() {
+    let TestSetup { lz_receiver_dispatcher, .. } = setup();
+
+    let is_allowed = lz_receiver_dispatcher.allow_initialize_path(MOCK_ORIGIN());
+    assert(is_allowed, 'Path should be allowed');
+
+    let unknown_origin = Origin {
+        src_eid: UNKNOWN_EID, sender: Bytes32 { value: 0x999 }, nonce: 1,
+    };
+    let is_not_allowed = lz_receiver_dispatcher.allow_initialize_path(unknown_origin);
+    assert(!is_not_allowed, 'Unknown path not allowed');
+}
+
+#[test]
+fn test_next_nonce() {
+    let TestSetup { lz_receiver_dispatcher, .. } = setup();
+
+    let nonce = lz_receiver_dispatcher.next_nonce(SRC_EID, MOCK_PEER_1());
+    assert(nonce == 0, 'Default nonce should be 0');
+}
+
+// =============================== Test OApp Receiver =================================
+
+#[test]
+fn test_is_compose_msg_sender() {
+    let TestSetup { oapp_address, oapp_receiver_dispatcher, .. } = setup();
+
+    let is_sender = oapp_receiver_dispatcher
+        .is_compose_msg_sender(MOCK_ORIGIN(), MOCK_MESSAGE(), oapp_address);
+    assert(is_sender, 'Contract should be valid sender');
+
+    let is_not_sender = oapp_receiver_dispatcher
+        .is_compose_msg_sender(MOCK_ORIGIN(), MOCK_MESSAGE(), NON_OWNER);
+    assert(!is_not_sender, 'Other address not valid sender');
+}
+
+// =============================== Test OApp Sender =================================
+
+// =============================== Test Quote =================================
+
+#[test]
+fn test_quote_success() {
+    let TestSetup { mock_oapp_core_dispatcher, .. } = setup();
+
+    // mock endpoint to return a 100 10 message fee
+    let expected_result = MessagingFee { native_fee: 100, lz_token_fee: 10 };
+
+    start_mock_call(FAKE_ENDPOINT, selector!("quote"), expected_result.clone());
+    let fee = mock_oapp_core_dispatcher
+        .test_quote(EID_1, MOCK_MESSAGE(), Default::default(), false);
+
+    assert(fee == expected_result, 'Fee should be 100 10');
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_quote_fails_unknown_peer() {
+    let TestSetup { mock_oapp_core_safe_dispatcher, .. } = setup();
+
+    let result = mock_oapp_core_safe_dispatcher
+        .test_quote(UNKNOWN_EID, MOCK_MESSAGE(), Default::default(), false);
+
+    assert_panic_with_error(result, err_no_peer(UNKNOWN_EID));
+}
+
+// =============================== Test Pay in Token =================================
+
+#[test]
+fn test_pay_in_token() {
+    // Should take the money into Oapp and approve the endpoint to spend it
+    let TestSetup { oapp_address, mock_oapp_core_dispatcher, stark_token_dispatcher, .. } = setup();
+    start_cheat_caller_address(stark_token_dispatcher.contract_address, OWNER);
+    stark_token_dispatcher.approve(oapp_address, 100);
+    stop_cheat_caller_address(stark_token_dispatcher.contract_address);
+    mock_oapp_core_dispatcher
+        .test_pay_in_token(
+            OWNER, FAKE_ENDPOINT, oapp_address, 100, stark_token_dispatcher.contract_address,
+        );
+
+    let stark_balance = stark_token_dispatcher.balance_of(oapp_address);
+    assert(stark_balance == 100, 'STARK token not taken');
+    let endpoint_allowance = stark_token_dispatcher.allowance(oapp_address, FAKE_ENDPOINT);
+    assert(endpoint_allowance == 100, 'STARK token not approved');
+}
+
+// =============================== Test Pay LZ Token =================================
+
+#[test]
+fn test_pay_lz_token() {
+    let TestSetup { oapp_address, mock_oapp_core_dispatcher, lz_token_dispatcher, .. } = setup();
+    start_cheat_caller_address(lz_token_dispatcher.contract_address, OWNER);
+    lz_token_dispatcher.approve(oapp_address, 100);
+    stop_cheat_caller_address(lz_token_dispatcher.contract_address);
+    mock_oapp_core_dispatcher.test_pay_lz_token(OWNER, FAKE_ENDPOINT, oapp_address, 100);
+
+    let lz_balance = lz_token_dispatcher.balance_of(oapp_address);
+    assert(lz_balance == 100, 'LZ token not taken');
+    let endpoint_allowance = lz_token_dispatcher.allowance(oapp_address, FAKE_ENDPOINT);
+    assert(endpoint_allowance == 100, 'LZ token not approved');
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_pay_lz_token_fails_not_enough_lz_token_allowance() {
+    let TestSetup { oapp_address, mock_oapp_core_safe_dispatcher, .. } = setup();
+    start_cheat_caller_address(oapp_address, OWNER);
+    let result = mock_oapp_core_safe_dispatcher
+        .test_pay_lz_token(OWNER, FAKE_ENDPOINT, oapp_address, 100);
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_error(result, err_not_enough_lz_token_allowance(100, 0));
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_pay_lz_token_fails_not_enough_lz_token() {
+    let TestSetup { oapp_address, mock_oapp_core_safe_dispatcher, .. } = setup();
+    start_cheat_caller_address(oapp_address, OWNER);
+    let result = mock_oapp_core_safe_dispatcher
+        .test_pay_lz_token(NON_OWNER, FAKE_ENDPOINT, oapp_address, 100);
+    stop_cheat_caller_address(oapp_address);
+    assert_panic_with_error(result, err_not_enough_lz_token(100, 0));
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_pay_lz_token_fails_lz_token_unavailable() {
+    // Deploy OApp with zero LZ token address
+    let stark_token = deploy_test_erc20(OWNER);
+    let oapp_address = deploy_test_oapp_core(FAKE_ENDPOINT, OWNER, stark_token);
+    // mock lz token to be unavailable
+    start_mock_call(FAKE_ENDPOINT, selector!("get_lz_token"), ZERO_ADDRESS);
+    let mock_oapp_core_safe_dispatcher = IMockOAppCoreSafeDispatcher {
+        contract_address: oapp_address,
+    };
+
+    let result = mock_oapp_core_safe_dispatcher
+        .test_pay_lz_token(OWNER, FAKE_ENDPOINT, oapp_address, 100);
+
+    assert_panic_with_error(result, err_lz_token_unavailable());
+}
+
+// =============================== Test Pay Native =================================
+
+#[test]
+fn test_pay_native() {
+    let TestSetup { oapp_address, mock_oapp_core_dispatcher, stark_token_dispatcher, .. } = setup();
+    start_cheat_caller_address(stark_token_dispatcher.contract_address, OWNER);
+    stark_token_dispatcher.approve(oapp_address, 100);
+    stop_cheat_caller_address(stark_token_dispatcher.contract_address);
+    mock_oapp_core_dispatcher.test_pay_native(OWNER, FAKE_ENDPOINT, oapp_address, 100);
+
+    let stark_balance = stark_token_dispatcher.balance_of(oapp_address);
+    assert(stark_balance == 100, 'STARK token not taken');
+    let endpoint_allowance = stark_token_dispatcher.allowance(oapp_address, FAKE_ENDPOINT);
+    assert(endpoint_allowance == 100, 'STARK token not approved');
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_pay_native_fails_not_enough_native_allowance() {
+    let TestSetup { oapp_address, mock_oapp_core_safe_dispatcher, .. } = setup();
+    start_cheat_caller_address(oapp_address, OWNER);
+    let result = mock_oapp_core_safe_dispatcher
+        .test_pay_native(OWNER, FAKE_ENDPOINT, oapp_address, 100);
+    stop_cheat_caller_address(oapp_address);
+    assert_panic_with_error(result, err_not_enough_native_allowance(100, 0));
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_pay_native_fails_not_enough_native() {
+    let TestSetup { oapp_address, mock_oapp_core_safe_dispatcher, .. } = setup();
+    start_cheat_caller_address(oapp_address, OWNER);
+    let result = mock_oapp_core_safe_dispatcher
+        .test_pay_native(NON_OWNER, FAKE_ENDPOINT, oapp_address, 100);
+    stop_cheat_caller_address(oapp_address);
+    assert_panic_with_error(result, err_not_enough_native(100, 0));
+}
+
+// =============================== Test LZ Send =================================
+
+#[test]
+fn test_lz_send_success() {
+    let TestSetup {
+        oapp_address, mock_oapp_core_dispatcher, stark_token_dispatcher, lz_token_dispatcher, ..,
+    } = setup();
+    let message = MOCK_MESSAGE();
+    let options = Default::default();
+    let fee = MessagingFee { native_fee: 100, lz_token_fee: 10 };
+    let refund_address = OWNER;
+
+    let expected_receipt = MessageReceipt {
+        guid: Bytes32 { value: 0x1 }, nonce: 0, payees: array![],
+    };
+
+    start_mock_call(FAKE_ENDPOINT, selector!("send"), expected_receipt.clone());
+
+    // Set up STARK token approval
+    start_cheat_caller_address(stark_token_dispatcher.contract_address, OWNER);
+    stark_token_dispatcher.approve(oapp_address, 100);
+    stop_cheat_caller_address(stark_token_dispatcher.contract_address);
+
+    // Set up LZ token approval
+    start_cheat_caller_address(lz_token_dispatcher.contract_address, OWNER);
+    lz_token_dispatcher.approve(oapp_address, 10);
+    stop_cheat_caller_address(lz_token_dispatcher.contract_address);
+
+    // Now make the lz_send call with OWNER as caller to the OApp
+    start_cheat_caller_address(oapp_address, OWNER);
+    let receipt = mock_oapp_core_dispatcher
+        .test_lz_send(EID_1, message, options, fee, refund_address);
+    stop_cheat_caller_address(oapp_address);
+
+    assert(receipt == expected_receipt, 'Receipt should be correct');
+
+    // Check that money was taken from the caller and allowance was set to endpoint
+    // in this mock the endpoint doesn't do anything so the allowance will still be the
+    // given values
+    let stark_balance = stark_token_dispatcher.balance_of(OWNER);
+    assert(stark_balance == INITIAL_SUPPLY - 100, 'STARK token not taken');
+    let oapp_stark_balance = stark_token_dispatcher.balance_of(oapp_address);
+    assert(oapp_stark_balance == 100, 'STARK token not received');
+    let endpoint_allowance = stark_token_dispatcher.allowance(oapp_address, FAKE_ENDPOINT);
+    assert(endpoint_allowance == 100, 'STARK token not approved');
+
+    let lz_balance = lz_token_dispatcher.balance_of(OWNER);
+    assert(lz_balance == INITIAL_SUPPLY - 10, 'LZ token not taken');
+    let oapp_lz_balance = lz_token_dispatcher.balance_of(oapp_address);
+    assert(oapp_lz_balance == 10, 'LZ token not received');
+    let endpoint_allowance = lz_token_dispatcher.allowance(oapp_address, FAKE_ENDPOINT);
+    assert(endpoint_allowance == 10, 'LZ token not approved');
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_send_fails_not_peer() {
+    let TestSetup { mock_oapp_core_safe_dispatcher, .. } = setup();
+
+    let result = mock_oapp_core_safe_dispatcher
+        .test_lz_send(
+            UNKNOWN_EID,
+            MOCK_MESSAGE(),
+            Default::default(),
+            MessagingFee { native_fee: 0, lz_token_fee: 0 },
+            OWNER,
+        );
+
+    assert_panic_with_error(result, err_no_peer(UNKNOWN_EID));
+}
+
+
+// =============================== Essential LZ Receive Tests =================================
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_receive_fails_wrong_caller() {
+    let TestSetup { oapp_address, lz_receiver_safe_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, NON_OWNER);
+    let result = lz_receiver_safe_dispatcher
+        .lz_receive(MOCK_ORIGIN(), MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_error(result, err_only_endpoint(FAKE_ENDPOINT));
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_receive_fails_unregistered_peer() {
+    let TestSetup { oapp_address, lz_receiver_safe_dispatcher, .. } = setup();
+    let unknown_origin = Origin { src_eid: UNKNOWN_EID, sender: MOCK_PEER_1(), nonce: 1 };
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+    let result = lz_receiver_safe_dispatcher
+        .lz_receive(unknown_origin, MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+
+    assert_panic_with_error(result, err_no_peer(UNKNOWN_EID));
+}
+
+#[test]
+fn test_lz_receive_multi_chain() {
+    let TestSetup { oapp_address, oapp_dispatcher, lz_receiver_dispatcher, .. } = setup();
+
+    // Set up peers for different chains
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    oapp_dispatcher.set_peer(EID_2, MOCK_PEER_2());
+    stop_cheat_caller_address(oapp_address);
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+
+    let origin_chain1 = Origin { src_eid: EID_1, sender: MOCK_PEER_1(), nonce: 1 };
+    let origin_chain2 = Origin { src_eid: EID_2, sender: MOCK_PEER_2(), nonce: 1 };
+
+    lz_receiver_dispatcher
+        .lz_receive(origin_chain1, MOCK_GUID(), "chain1 msg", EXECUTOR, 0, Default::default());
+    lz_receiver_dispatcher
+        .lz_receive(origin_chain2, MOCK_GUID(), "chain2 msg", EXECUTOR, 0, Default::default());
+
+    stop_cheat_caller_address(oapp_address);
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_lz_receive_fails_wrong_sender_for_chain() {
+    let TestSetup { oapp_address, oapp_dispatcher, lz_receiver_safe_dispatcher, .. } = setup();
+
+    // Set up peer for EID_1
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_peer(EID_1, MOCK_PEER_1());
+    stop_cheat_caller_address(oapp_address);
+
+    // Try wrong sender for EID_1
+    let invalid_origin = Origin { src_eid: EID_1, sender: MOCK_PEER_2(), nonce: 1 };
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+    let result = lz_receiver_safe_dispatcher
+        .lz_receive(invalid_origin, MOCK_GUID(), MOCK_MESSAGE(), EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+    assert_panic_with_error(result, err_only_peer(EID_1, MOCK_PEER_2()));
+}
+
+#[test]
+fn test_lz_receive_empty_message() {
+    let TestSetup { oapp_address, lz_receiver_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, FAKE_ENDPOINT);
+    lz_receiver_dispatcher
+        .lz_receive(MOCK_ORIGIN(), MOCK_GUID(), "", EXECUTOR, 0, Default::default());
+    stop_cheat_caller_address(oapp_address);
+}
+
+// =============================== Test Set Delegate =================================
+
+#[test]
+fn test_set_delegate() {
+    let TestSetup { oapp_address, oapp_dispatcher, .. } = setup();
+
+    start_mock_call(FAKE_ENDPOINT, selector!("set_delegate"), OWNER);
+
+    start_cheat_caller_address(oapp_address, OWNER);
+    oapp_dispatcher.set_delegate(OWNER);
+    stop_cheat_caller_address(oapp_address);
+}
+
+#[test]
+#[feature("safe_dispatcher")]
+fn test_set_delegate_fails_not_owner() {
+    let TestSetup { oapp_address, oapp_safe_dispatcher, .. } = setup();
+
+    start_cheat_caller_address(oapp_address, NON_OWNER);
+    let result = oapp_safe_dispatcher.set_delegate(OWNER);
+    stop_cheat_caller_address(oapp_address);
+    assert_panic_with_felt_error(result, OwnableComponent::Errors::NOT_OWNER);
+}