Test token hooks (#1198)

* erc20 hooks tests * add erc721 tests * use trait helper * simplify mocks * fmt * Update packages/test_common/src/mocks/erc20.cairo Co-authored-by: immrsd <[email protected]> * remove unnecesary cheat * add erc1155 tests * rename mocks --------- Co-authored-by: immrsd <[email protected]>
OpenZeppelin · Nov 7, 2024 · 5fec9a5 · 5fec9a5
1 parent 2eda08b
commit 5fec9a5
Show file tree

Hide file tree

Showing 6 changed files with 487 additions and 6 deletions.
diff --git a/packages/test_common/src/mocks/erc1155.cairo b/packages/test_common/src/mocks/erc1155.cairo
@@ -102,6 +102,103 @@ pub mod SnakeERC1155Mock {
     }
 }
 
+/// Similar to `SnakeERC1155Mock`, but emits events for `before_update` and `after_update` hooks.
+/// This is used to test that the hooks are called with the correct arguments.
+#[starknet::contract]
+pub mod SnakeERC1155MockWithHooks {
+    use openzeppelin_introspection::src5::SRC5Component;
+    use openzeppelin_token::erc1155::{ERC1155Component};
+    use starknet::ContractAddress;
+
+    component!(path: ERC1155Component, storage: erc1155, event: ERC1155Event);
+    component!(path: SRC5Component, storage: src5, event: SRC5Event);
+
+    // ERC1155
+    #[abi(embed_v0)]
+    impl ERC1155Impl = ERC1155Component::ERC1155Impl<ContractState>;
+    #[abi(embed_v0)]
+    impl ERC1155MetadataURIImpl =
+        ERC1155Component::ERC1155MetadataURIImpl<ContractState>;
+    impl ERC1155InternalImpl = ERC1155Component::InternalImpl<ContractState>;
+
+    // SRC5
+    #[abi(embed_v0)]
+    impl SRC5Impl = SRC5Component::SRC5Impl<ContractState>;
+
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub erc1155: ERC1155Component::Storage,
+        #[substorage(v0)]
+        pub src5: SRC5Component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    pub enum Event {
+        #[flat]
+        ERC1155Event: ERC1155Component::Event,
+        #[flat]
+        SRC5Event: SRC5Component::Event,
+        BeforeUpdate: BeforeUpdate,
+        AfterUpdate: AfterUpdate
+    }
+
+    /// Event used to test that `before_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct BeforeUpdate {
+        pub from: ContractAddress,
+        pub to: ContractAddress,
+        pub token_ids: Span<u256>,
+        pub values: Span<u256>
+    }
+
+    /// Event used to test that `after_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct AfterUpdate {
+        pub from: ContractAddress,
+        pub to: ContractAddress,
+        pub token_ids: Span<u256>,
+        pub values: Span<u256>
+    }
+
+    #[constructor]
+    fn constructor(
+        ref self: ContractState,
+        base_uri: ByteArray,
+        recipient: ContractAddress,
+        token_id: u256,
+        value: u256
+    ) {
+        self.erc1155.initializer(base_uri);
+        self.erc1155.mint_with_acceptance_check(recipient, token_id, value, array![].span());
+    }
+
+    impl ERC1155HooksImpl of ERC1155Component::ERC1155HooksTrait<ContractState> {
+        fn before_update(
+            ref self: ERC1155Component::ComponentState<ContractState>,
+            from: ContractAddress,
+            to: ContractAddress,
+            token_ids: Span<u256>,
+            values: Span<u256>
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(BeforeUpdate { from, to, token_ids, values });
+        }
+
+        fn after_update(
+            ref self: ERC1155Component::ComponentState<ContractState>,
+            from: ContractAddress,
+            to: ContractAddress,
+            token_ids: Span<u256>,
+            values: Span<u256>
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(AfterUpdate { from, to, token_ids, values });
+        }
+    }
+}
+
 #[starknet::contract]
 pub mod DualCaseERC1155ReceiverMock {
     use openzeppelin_introspection::src5::SRC5Component;

diff --git a/packages/test_common/src/mocks/erc20.cairo b/packages/test_common/src/mocks/erc20.cairo
@@ -78,6 +78,87 @@ pub mod SnakeERC20Mock {
     }
 }
 
+/// Similar to `SnakeERC20Mock`, but emits events for `before_update` and `after_update` hooks.
+/// This is used to test that the hooks are called with the correct arguments.
+#[starknet::contract]
+pub mod SnakeERC20MockWithHooks {
+    use openzeppelin_token::erc20::ERC20Component;
+    use starknet::ContractAddress;
+
+    component!(path: ERC20Component, storage: erc20, event: ERC20Event);
+
+    #[abi(embed_v0)]
+    impl ERC20Impl = ERC20Component::ERC20Impl<ContractState>;
+    #[abi(embed_v0)]
+    impl ERC20MetadataImpl = ERC20Component::ERC20MetadataImpl<ContractState>;
+    impl InternalImpl = ERC20Component::InternalImpl<ContractState>;
+
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub erc20: ERC20Component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    pub enum Event {
+        #[flat]
+        ERC20Event: ERC20Component::Event,
+        BeforeUpdate: BeforeUpdate,
+        AfterUpdate: AfterUpdate
+    }
+
+    /// Event used to test that `before_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct BeforeUpdate {
+        pub from: ContractAddress,
+        pub recipient: ContractAddress,
+        pub amount: u256
+    }
+
+    /// Event used to test that `after_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct AfterUpdate {
+        pub from: ContractAddress,
+        pub recipient: ContractAddress,
+        pub amount: u256
+    }
+
+    #[constructor]
+    fn constructor(
+        ref self: ContractState,
+        name: ByteArray,
+        symbol: ByteArray,
+        initial_supply: u256,
+        recipient: ContractAddress
+    ) {
+        self.erc20.initializer(name, symbol);
+        self.erc20.mint(recipient, initial_supply);
+    }
+
+    impl ERC20HooksImpl of ERC20Component::ERC20HooksTrait<ContractState> {
+        fn before_update(
+            ref self: ERC20Component::ComponentState<ContractState>,
+            from: ContractAddress,
+            recipient: ContractAddress,
+            amount: u256
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(BeforeUpdate { from, recipient, amount });
+        }
+
+        fn after_update(
+            ref self: ERC20Component::ComponentState<ContractState>,
+            from: ContractAddress,
+            recipient: ContractAddress,
+            amount: u256
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(AfterUpdate { from, recipient, amount });
+        }
+    }
+}
+
 #[starknet::contract]
 pub mod DualCaseERC20PermitMock {
     use openzeppelin_token::erc20::{ERC20Component, ERC20HooksEmptyImpl};

diff --git a/packages/test_common/src/mocks/erc721.cairo b/packages/test_common/src/mocks/erc721.cairo
@@ -107,6 +107,99 @@ pub mod SnakeERC721Mock {
     }
 }
 
+/// Similar as `SnakeERC721Mock`, but emits events for `before_update` and `after_update` hooks.
+/// This is used to test that the hooks are called with the correct arguments.
+#[starknet::contract]
+pub mod SnakeERC721MockWithHooks {
+    use openzeppelin_introspection::src5::SRC5Component;
+    use openzeppelin_token::erc721::ERC721Component;
+    use starknet::ContractAddress;
+
+    component!(path: ERC721Component, storage: erc721, event: ERC721Event);
+    component!(path: SRC5Component, storage: src5, event: SRC5Event);
+
+    // ERC721
+    #[abi(embed_v0)]
+    impl ERC721Impl = ERC721Component::ERC721Impl<ContractState>;
+    #[abi(embed_v0)]
+    impl ERC721MetadataImpl = ERC721Component::ERC721MetadataImpl<ContractState>;
+    impl ERC721InternalImpl = ERC721Component::InternalImpl<ContractState>;
+
+    // SRC5
+    #[abi(embed_v0)]
+    impl SRC5Impl = SRC5Component::SRC5Impl<ContractState>;
+
+    #[storage]
+    pub struct Storage {
+        #[substorage(v0)]
+        pub erc721: ERC721Component::Storage,
+        #[substorage(v0)]
+        pub src5: SRC5Component::Storage
+    }
+
+    #[event]
+    #[derive(Drop, starknet::Event)]
+    pub enum Event {
+        #[flat]
+        ERC721Event: ERC721Component::Event,
+        #[flat]
+        SRC5Event: SRC5Component::Event,
+        BeforeUpdate: BeforeUpdate,
+        AfterUpdate: AfterUpdate
+    }
+
+    /// Event used to test that `before_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct BeforeUpdate {
+        pub to: ContractAddress,
+        pub token_id: u256,
+        pub auth: ContractAddress
+    }
+
+    /// Event used to test that `after_update` hook is called.
+    #[derive(Drop, PartialEq, starknet::Event)]
+    pub struct AfterUpdate {
+        pub to: ContractAddress,
+        pub token_id: u256,
+        pub auth: ContractAddress
+    }
+
+    #[constructor]
+    fn constructor(
+        ref self: ContractState,
+        name: ByteArray,
+        symbol: ByteArray,
+        base_uri: ByteArray,
+        recipient: ContractAddress,
+        token_id: u256
+    ) {
+        self.erc721.initializer(name, symbol, base_uri);
+        self.erc721.mint(recipient, token_id);
+    }
+
+    impl ERC721HooksImpl of ERC721Component::ERC721HooksTrait<ContractState> {
+        fn before_update(
+            ref self: ERC721Component::ComponentState<ContractState>,
+            to: ContractAddress,
+            token_id: u256,
+            auth: ContractAddress
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(BeforeUpdate { to, token_id, auth });
+        }
+
+        fn after_update(
+            ref self: ERC721Component::ComponentState<ContractState>,
+            to: ContractAddress,
+            token_id: u256,
+            auth: ContractAddress
+        ) {
+            let mut contract_state = self.get_contract_mut();
+            contract_state.emit(AfterUpdate { to, token_id, auth });
+        }
+    }
+}
+
 #[starknet::contract]
 pub mod DualCaseERC721ReceiverMock {
     use openzeppelin_introspection::src5::SRC5Component;