🚀 add team allocation check

contracts/src/tokens/memecoin.cairo

@@ -1,11 +1,11 @@
 //! `UnruggableMemecoin` is an ERC20 token has additional features to prevent rug pulls.
 use starknet::ContractAddress;
 
-
 #[starknet::contract]
 mod UnruggableMemecoin {
     use integer::BoundedInt;
     use openzeppelin::access::ownable::OwnableComponent;
+    use openzeppelin::access::ownable::interface::IOwnable;
     use openzeppelin::access::ownable::ownable::OwnableComponent::InternalTrait;
     use openzeppelin::token::erc20::ERC20Component;
     use starknet::{ContractAddress, get_caller_address};
@@ -94,12 +94,26 @@ mod UnruggableMemecoin {
         fn launch_memecoin(ref self: ContractState) {
             // Checks: Only the owner can launch the memecoin.
             self.ownable.assert_only_owner();
+            // Checks team allocation, should be less than 10% of the total supply.
+            self.check_team_allocation(0);
         // Effects.
 
         // Interactions.
+
+        }
+
+        /// Compute the maximum team allocation
+        /// total supply * MAX_SUPPLY_PERCENTAGE_TEAM_ALLOCATION / 100
+        /// # RETURNS
+        /// - amount of tokens
+        fn get_max_team_allocation(self: @ContractState) -> u256 {
+            self.erc20.ERC20_total_supply.read()
+                * MAX_SUPPLY_PERCENTAGE_TEAM_ALLOCATION.into()
+                / 100
         }
     }
 
+
     #[abi(embed_v0)]
     impl SnakeEntrypoints of IUnruggableMemecoinSnake<ContractState> {
         // ************************************
@@ -120,6 +134,9 @@ mod UnruggableMemecoin {
         }
 
         fn transfer(ref self: ContractState, recipient: ContractAddress, amount: u256) -> bool {
+            if (recipient == self.ownable.owner()) {
+                self.check_team_allocation(amount);
+            }
             let sender = get_caller_address();
             self._transfer(sender, recipient, amount);
             true
@@ -131,6 +148,9 @@ mod UnruggableMemecoin {
             recipient: ContractAddress,
             amount: u256
         ) -> bool {
+            if (recipient == self.ownable.owner()) {
+                self.check_team_allocation(amount);
+            }
             let caller = get_caller_address();
             self.erc20._spend_allowance(sender, caller, amount);
             self.erc20._transfer(sender, recipient, amount);
@@ -179,5 +199,12 @@ mod UnruggableMemecoin {
         ) {
             self.erc20._transfer(sender, recipient, amount);
         }
+        /// Check if the team allocation cap is reached
+        fn check_team_allocation(self: @ContractState, amount: u256) {
+            assert(
+                self.get_max_team_allocation() >= (self.balance_of(self.ownable.owner()) + amount),
+                'Team allocation cap reached'
+            );
+        }
     }
 }

contracts/tests/test_unruggable_memecoin.cairo

@@ -1,7 +1,6 @@
-use core::debug::PrintTrait;
 use core::traits::Into;
 use openzeppelin::token::erc20::interface::IERC20;
-use snforge_std::{declare, ContractClassTrait, start_prank, stop_prank};
+use snforge_std::{declare, ContractClassTrait, start_prank, stop_prank, CheatTarget};
 use starknet::{ContractAddress, contract_address_const};
 
 use unruggable::tokens::interface::{
@@ -330,3 +329,117 @@ mod memecoin_entrypoints {
         memecoin.launch_memecoin();
     }
 }
+
+
+#[test]
+fn test_team_alloc() {
+    let owner = contract_address_const::<1>();
+    let recipient = contract_address_const::<2>();
+    let initial_supply = 1000.into();
+    let contract_address = deploy_contract(
+        owner, recipient, 'UnruggableMemecoin', 'URM', initial_supply
+    );
+    // set owner as caller
+    start_prank(CheatTarget::One(contract_address), owner);
+
+    // make transfer 
+    let safe_dispatcher = IUnruggableMemecoinDispatcher { contract_address };
+    let max_alloc: u256 = safe_dispatcher.get_max_team_allocation();
+
+    assert(max_alloc == 100, 'Invalid max allocation');
+}
+#[test]
+#[should_panic(expected: ('Team allocation cap reached',))]
+fn test_fail_transfer_team_alloc_reached() {
+    let owner = contract_address_const::<1>();
+    let recipient = contract_address_const::<2>();
+    let initial_supply = 1000.into();
+    let contract_address = deploy_contract(
+        owner, recipient, 'UnruggableMemecoin', 'URM', initial_supply
+    );
+
+    // make transfer 
+    let safe_dispatcher = IUnruggableMemecoinDispatcher { contract_address };
+    // Check initial balance. Should be equal to initial supply.
+    let mut balance = safe_dispatcher.balance_of(owner);
+    assert(balance == 0, 'Invalid owner balance');
+    balance = safe_dispatcher.balance_of(recipient);
+    assert(balance == initial_supply, 'Invalid recipient balance');
+
+    // set recipient as caller
+    start_prank(CheatTarget::One(contract_address), recipient);
+    safe_dispatcher.transfer(owner, initial_supply);
+}
+
+#[test]
+fn test_transfer_team_alloc_not_reached() {
+    let owner = contract_address_const::<1>();
+    let recipient = contract_address_const::<2>();
+    let initial_supply = 1000.into();
+    let contract_address = deploy_contract(
+        owner, recipient, 'UnruggableMemecoin', 'URM', initial_supply
+    );
+
+    // make transfer 
+    let safe_dispatcher = IUnruggableMemecoinDispatcher { contract_address };
+    // Check initial balance. Should be equal to initial supply.
+    let mut balance = safe_dispatcher.balance_of(owner);
+    assert(balance == 0, 'Invalid owner balance');
+    balance = safe_dispatcher.balance_of(recipient);
+    assert(balance == initial_supply, 'Invalid recipient balance');
+
+    // set recipient as caller
+    start_prank(CheatTarget::One(contract_address), recipient);
+    safe_dispatcher.transfer(owner, 100);
+}
+
+#[test]
+#[should_panic(expected: ('Team allocation cap reached',))]
+fn test_fail_transfer_from_team_alloc_reached() {
+    let owner = contract_address_const::<1>();
+    let recipient = contract_address_const::<2>();
+    let initial_supply = 1000.into();
+    let contract_address = deploy_contract(
+        owner, recipient, 'UnruggableMemecoin', 'URM', initial_supply
+    );
+
+    // make transfer 
+    let safe_dispatcher = IUnruggableMemecoinDispatcher { contract_address };
+    // Check initial balance. Should be equal to initial supply.
+    let mut balance = safe_dispatcher.balance_of(owner);
+    assert(balance == 0, 'Invalid owner balance');
+    balance = safe_dispatcher.balance_of(recipient);
+    assert(balance == initial_supply, 'Invalid recipient balance');
+
+    // set recipient as caller
+    start_prank(CheatTarget::One(contract_address), recipient);
+    safe_dispatcher.approve(owner, initial_supply);
+    // set owner as caller to transfer from recipient
+    start_prank(CheatTarget::One(contract_address), owner);
+    safe_dispatcher.transfer_from(recipient, owner, initial_supply);
+}
+
+#[test]
+fn test_transfer_from_team_alloc_not_reached() {
+    let owner = contract_address_const::<1>();
+    let recipient = contract_address_const::<2>();
+    let initial_supply = 1000.into();
+    let contract_address = deploy_contract(
+        owner, recipient, 'UnruggableMemecoin', 'URM', initial_supply
+    );
+
+    // make transfer 
+    let safe_dispatcher = IUnruggableMemecoinDispatcher { contract_address };
+    // Check initial balance. Should be equal to initial supply.
+    let mut balance = safe_dispatcher.balance_of(owner);
+    assert(balance == 0, 'Invalid owner balance');
+    balance = safe_dispatcher.balance_of(recipient);
+    assert(balance == initial_supply, 'Invalid recipient balance');
+
+    // set recipient as caller
+    start_prank(CheatTarget::One(contract_address), recipient);
+    safe_dispatcher.approve(owner, 100);
+    // set owner as caller to transfer from recipient
+    start_prank(CheatTarget::One(contract_address), owner);
+    safe_dispatcher.transfer_from(recipient, owner, 100);
+}