iwa 0.0.10__py3-none-any.whl → 0.0.12__py3-none-any.whl

iwa/plugins/olas/service_manager/staking.py

@@ -1,4 +1,48 @@
-"""Staking manager mixin."""
+"""Staking manager mixin for OLAS service staking operations.
+
+OLAS Token Flow Overview
+========================
+
+For OLAS token-bonded services (e.g., Expert 7 MM requiring 10,000 OLAS total),
+the tokens flow through multiple stages:
+
+1. CREATE SERVICE
+   - Service is registered on-chain with bond parameters
+   - Service Owner approves Token Utility to spend OLAS (2 × bond)
+   - NO OLAS tokens move yet
+
+2. ACTIVATION (min_staking_deposit = 5,000 OLAS for 10k contract)
+   - Service Owner approves Token Utility for the security deposit
+   - TX sends 1 wei native value (not 5k OLAS!)
+   - Token Utility internally calls transferFrom() to move 5k OLAS
+   - 5k OLAS moves: Service Owner → Token Utility
+
+3. REGISTRATION (agent_bond = 5,000 OLAS for 10k contract)
+   - Service Owner approves Token Utility for the bond amount
+   - TX sends 1 wei native value per agent (not 5k OLAS!)
+   - Token Utility internally calls transferFrom() to move 5k OLAS
+   - 5k OLAS moves: Service Owner → Token Utility
+
+4. DEPLOY
+   - Creates the Safe multisig for the service
+   - NO OLAS tokens move
+
+5. STAKE (this module) ★
+   - Only the Service NFT is approved to the staking contract
+   - NO OLAS tokens move in this transaction!
+   - The staking contract reads the deposited amounts from Token Utility
+   - Service Registry L2 token (NFT) moves: Owner → Staking Contract
+
+Key Insight:
+    At stake time, the Service Owner's OLAS balance is 0 (all 10k was deposited
+    during activation + registration). This is correct! The staking contract
+    pulls position data from the Token Utility, not from the owner's wallet.
+
+Contract Addresses (Gnosis):
+    - Token Utility: 0xa45E...8eD8
+    - Service Registry L2: 0x9338...55fD
+    - OLAS Token: 0xcE11...d9f
+"""
 
 from datetime import datetime, timezone
 from typing import Optional
@@ -6,7 +50,6 @@ from typing import Optional
 from loguru import logger
 from web3 import Web3
 
-from iwa.core.contracts.erc20 import ERC20Contract
 from iwa.core.types import EthereumAddress
 from iwa.core.utils import get_tx_hash
 from iwa.plugins.olas.contracts.staking import StakingContract, StakingState
@@ -14,7 +57,21 @@ from iwa.plugins.olas.models import StakingStatus
 
 
 class StakingManagerMixin:
-    """Mixin for staking operations."""
+    """Mixin for staking operations on OLAS services.
+
+    This mixin handles the final step of the service lifecycle: staking a
+    deployed service into a staking contract to earn OLAS rewards.
+
+    Important: By the time stake() is called, all OLAS tokens have already
+    been deposited to the Token Utility during activation and registration.
+    The stake transaction only transfers the Service NFT, not OLAS tokens.
+
+    Staking Requirements:
+        - Service must be in DEPLOYED state
+        - Service must be created with OLAS token (not native currency)
+        - Staking contract must have available slots
+        - Service token must match staking contract's required token
+    """
 
     def get_staking_status(self) -> Optional[StakingStatus]:
         """Get comprehensive staking status for the active service.
@@ -154,20 +211,28 @@ class StakingManagerMixin:
     def stake(self, staking_contract) -> bool:
         """Stake the service in a staking contract.
 
-        Token Flow:
-            The total OLAS required is split 50/50 between deposit and bond:
-            - minStakingDeposit: Transferred to staking contract during this call
-            - agentBond: Already in Token Utility from service registration
-
-            Example for Hobbyist 1 (100 OLAS total):
-            - minStakingDeposit: 50 OLAS (from master account -> staking contract)
-            - agentBond: 50 OLAS (already in Token Utility)
-
-        Requirements:
-            - Service must be in DEPLOYED state
-            - Service must be created with OLAS token (not native currency)
-            - Master account must have >= minStakingDeposit OLAS tokens
-            - Staking contract must have available slots
+        This is the final step after create → activate → register → deploy.
+        At this point, all OLAS tokens are already in the Token Utility.
+
+        Token Flow at Stake Time:
+            ┌─────────────────────────────────────────────────────────────────┐
+            │  What Moves:                                                    │
+            │    • Service NFT (ERC-721): Owner → Staking Contract            │
+            │                                                                 │
+            │  What does NOT Move:                                            │
+            │    • OLAS tokens - already in Token Utility from earlier steps  │
+            └─────────────────────────────────────────────────────────────────┘
+
+        Why no OLAS transfer?
+            The staking contract reads the service's bond/deposit from the
+            Token Utility contract. It doesn't need a new transfer - it just
+            verifies the amounts are sufficient and locks the service.
+
+        Process:
+            1. Validate requirements (state, token, slots)
+            2. Approve Service NFT to staking contract
+            3. Call stake(serviceId) on staking contract
+            4. Verify ServiceStaked event
 
         Args:
             staking_contract: StakingContract instance to stake in.
@@ -194,12 +259,12 @@ class StakingManagerMixin:
             f"[STAKE] Min deposit required: {min_deposit} wei ({min_deposit / 1e18:.2f} OLAS)"
         )
 
-        # 2. Approve Tokens
-        logger.info("[STAKE] Step 2: Approving tokens...")
-        if not self._approve_staking_tokens(staking_contract, min_deposit):
-            logger.error("[STAKE] Step 2 FAILED: Token approval failed")
+        # 2. Approve Service NFT
+        logger.info("[STAKE] Step 2: Approving service NFT...")
+        if not self._approve_staking_tokens(staking_contract):
+            logger.error("[STAKE] Step 2 FAILED: NFT approval failed")
             return False
-        logger.info("[STAKE] Step 2 OK: Tokens approved")
+        logger.info("[STAKE] Step 2 OK: Service NFT approved")
 
         # 3. Execute Stake Transaction
         logger.info("[STAKE] Step 3: Executing stake transaction...")
@@ -214,7 +279,27 @@ class StakingManagerMixin:
         return result
 
     def _check_stake_requirements(self, staking_contract) -> Optional[dict]:
-        """Validate all conditions required for staking."""
+        """Validate all conditions required for staking.
+
+        Checks performed:
+            1. Service State: Must be DEPLOYED (multisig created)
+            2. Token Match: Service token == Staking contract's staking_token
+            3. Agent Bond: Logged (may show 1 wei on-chain, this is normal)
+            4. Available Slots: Contract must have free slots
+
+        Note on OLAS Balance:
+            We do NOT check owner's OLAS balance here. By this point:
+            - 5k OLAS was transferred during activation (to Token Utility)
+            - 5k OLAS was transferred during registration (to Token Utility)
+            - Owner's OLAS balance is 0, and that's correct!
+
+        Args:
+            staking_contract: StakingContract to validate against.
+
+        Returns:
+            Dict with {min_deposit, staking_token} if valid, None otherwise.
+
+        """
         from iwa.plugins.olas.contracts.service import ServiceState
 
         logger.debug("[STAKE] Fetching contract requirements...")
@@ -291,25 +376,36 @@ class StakingManagerMixin:
             return None
         logger.debug("[STAKE] OK: Slots available")
 
-        # Check OLAS balance
-        logger.debug("[STAKE] Checking master OLAS balance...")
-        erc20_contract = ERC20Contract(staking_token)
-        master_balance = erc20_contract.balance_of_wei(self.wallet.master_account.address)
-        logger.info(
-            f"[STAKE] Master OLAS balance: {master_balance} wei "
-            f"({master_balance / 1e18:.2f} OLAS, need {min_deposit / 1e18:.2f} OLAS)"
-        )
-
-        if master_balance < min_deposit:
-            logger.error(f"[STAKE] FAIL: Insufficient balance ({master_balance} < {min_deposit})")
-            return None
-        logger.debug("[STAKE] OK: Sufficient balance")
+        # NOTE: We don't check OLAS balance here because OLAS was already
+        # deposited to the Token Utility during activation (min_staking_deposit)
+        # and registration (agent_bond). The staking contract pulls from there.
+        logger.debug("[STAKE] OLAS already deposited to Token Utility during activation/registration")
 
         return {"min_deposit": min_deposit, "staking_token": staking_token}
 
-    def _approve_staking_tokens(self, staking_contract, min_deposit: int) -> bool:
-        """Approve both the service NFT and OLAS tokens for staking."""
-        # Approve service NFT
+    def _approve_staking_tokens(self, staking_contract) -> bool:
+        """Approve the Service NFT for transfer to the staking contract.
+
+        What This Does:
+            Calls approve(stakingContract, serviceId) on the Service Registry L2.
+            This allows the staking contract to transferFrom the NFT.
+
+        What This Does NOT Do:
+            - Does NOT approve OLAS tokens (they're already in Token Utility)
+            - Does NOT transfer any tokens (that happens in _execute_stake_transaction)
+
+        Token/NFT Movement:
+            BEFORE: Owner has NFT, staking contract has no approval
+            AFTER:  Owner has NFT, staking contract is approved to take it
+
+        Who Signs:
+            Master account (must be service owner)
+
+        Returns:
+            True if approval succeeded, False otherwise.
+
+        """
+        # Approve service NFT - this is an ERC-721 approval, not ERC-20
         logger.debug("[STAKE] Approving service NFT for staking contract...")
         approve_tx = self.registry.prepare_approve_tx(
             from_address=self.wallet.master_account.address,
@@ -330,36 +426,29 @@ class StakingManagerMixin:
 
         tx_hash = get_tx_hash(receipt)
         logger.info(f"[STAKE] Service NFT approved: {tx_hash}")
+        return True
 
-        # Approve OLAS tokens
-        logger.debug(f"[STAKE] Approving OLAS tokens ({min_deposit} wei)...")
-        reqs = staking_contract.get_requirements()
-        staking_token = Web3.to_checksum_address(reqs["staking_token"])
-        erc20_contract = ERC20Contract(staking_token)
+    def _execute_stake_transaction(self, staking_contract) -> bool:
+        """Execute the actual stake transaction on the staking contract.
 
-        olas_approve_tx = erc20_contract.prepare_approve_tx(
-            from_address=self.wallet.master_account.address,
-            spender=staking_contract.address,
-            amount_wei=min_deposit,
-        )
+        What Happens Internally:
+            1. Staking contract calls transferFrom to take the Service NFT
+            2. Staking contract reads bond/deposit from Token Utility
+            3. Staking contract records the service as staked
+            4. ServiceStaked event is emitted
 
-        success, receipt = self.wallet.sign_and_send_transaction(
-            transaction=olas_approve_tx,
-            signer_address_or_tag=self.wallet.master_account.address,
-            chain_name=self.chain_name,
-            tags=["olas_approve_olas_token"],
-        )
+        Token Movement:
+            - Service NFT: Owner → Staking Contract (via transferFrom)
+            - OLAS tokens: None! Already in Token Utility
 
-        if not success:
-            logger.error("[STAKE] FAIL: OLAS token approval failed")
-            return False
+        Why No OLAS Transfer?
+            The staking contract calls ServiceRegistryTokenUtility.getOperatorBalance()
+            to verify the deposited amounts. It doesn't need a new transfer.
 
-        tx_hash = get_tx_hash(receipt)
-        logger.info(f"[STAKE] OLAS tokens approved: {tx_hash}")
-        return True
+        Returns:
+            True if stake succeeded and ServiceStaked event was found.
 
-    def _execute_stake_transaction(self, staking_contract) -> bool:
-        """Send the stake transaction and verify the result."""
+        """
         logger.debug("[STAKE] Preparing stake transaction...")
         stake_tx = staking_contract.prepare_stake_tx(
             from_address=self.wallet.master_account.address,

iwa/plugins/olas/tests/test_olas_contracts.py

@@ -118,6 +118,8 @@ def test_service_manager_complex_registration(mock_erc20_cls, mock_wallet):
         manager.wallet.transfer_service.approve_erc20.return_value = True
         manager.registry.extract_events.return_value = [{"name": "RegisterInstance"}]
         mock_erc20_cls.return_value.balance_of_wei.return_value = 1000
+        # Fix: Mock allowance to return an int, not MagicMock
+        manager.wallet.transfer_service.get_erc20_allowance.return_value = 0
         assert manager.register_agent(VALID_ADDR, 100) is True
 
     # deploy successes
@@ -160,10 +162,12 @@ def test_service_manager_config_edges(mock_wallet):
     """Test ServiceManager configuration and initialization edge cases."""
     with patch("iwa.plugins.olas.service_manager.Config") as mock_cfg_cls:
         mock_cfg = mock_cfg_cls.return_value
-        mock_cfg.plugins = {"olas": MagicMock()}
-        mock_cfg.plugins["olas"].get_service.return_value = Service(
+        olas_mock = MagicMock()
+        olas_mock.get_service.return_value = Service(
             service_name="t", chain_name="gnosis", service_id=1, agent_ids=[1]
         )
+        # Ensure plugins.get("olas") returns our mock
+        mock_cfg.plugins = {"olas": olas_mock}
         with patch("iwa.plugins.olas.service_manager.ChainInterfaces"):
             # hits 56
             with patch(

iwa/plugins/olas/tests/test_service_lifecycle.py

@@ -95,10 +95,7 @@ def test_sm_stake_fail(sm):
     sm.transfer_service.approve_erc20.return_value = True
     sm.wallet.sign_and_send_transaction.return_value = (False, None)
 
-    with (
-        patch("iwa.plugins.olas.contracts.staking.StakingContract") as mock_stk_cls,
-        patch("iwa.plugins.olas.service_manager.staking.ERC20Contract"),
-    ):
+    with patch("iwa.plugins.olas.contracts.staking.StakingContract") as mock_stk_cls:
         mock_stk = mock_stk_cls.return_value
         mock_stk.get_service_info.return_value = {"staking_state": 1}
         mock_stk.staking_token_address = addr

iwa/plugins/olas/tests/test_service_manager.py

@@ -100,11 +100,6 @@ def mock_chain_interfaces():
         yield mock
 
 
-@pytest.fixture
-def mock_erc20_contract():
-    """Mock ERC20 contract fixture."""
-    with patch("iwa.plugins.olas.service_manager.staking.ERC20Contract") as mock:
-        yield mock
 
 
 @pytest.fixture
@@ -114,7 +109,7 @@ def service_manager(
     mock_registry,
     mock_manager_contract,
     mock_chain_interfaces,
-    mock_erc20_contract,
+
     mock_olas_config,
     mock_service,
 ):
@@ -174,12 +169,11 @@ def test_create_no_event(service_manager, mock_wallet):
     mock_wallet.sign_and_send_transaction.return_value = (True, {})
     service_manager.registry.extract_events.return_value = []
 
-    with patch("iwa.plugins.olas.service_manager.staking.ERC20Contract"):  # Mock ERC20
-        res = service_manager.create(
-            token_address_or_tag="0x1111111111111111111111111111111111111111"
-        )
-        # create() finds no ID, logs error, returns None for service_id.
-        assert res is None
+    res = service_manager.create(
+        token_address_or_tag="0x1111111111111111111111111111111111111111"
+    )
+    # create() finds no ID, logs error, returns None for service_id.
+    assert res is None
 
 
 def test_activate_registration_success(service_manager, mock_wallet):
@@ -193,7 +187,7 @@ def test_activate_registration_success(service_manager, mock_wallet):
 
     # Mock balance/allowance for the new check
     mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     assert service_manager.activate_registration() is True
@@ -287,18 +281,15 @@ def test_stake_success(service_manager, mock_wallet):
         "security_deposit": 50000000000000000000,
     }
 
-    with patch("iwa.plugins.olas.service_manager.staking.ERC20Contract") as mock_erc20:
-        mock_erc20.return_value.balance_of_wei.return_value = 100000000000000000000  # 100 OLAS
-
-        mock_wallet.sign_and_send_transaction.return_value = (True, {})
-        staking_contract.extract_events.return_value = [{"name": "ServiceStaked"}]
-        staking_contract.get_staking_state.return_value = StakingState.STAKED
+    mock_wallet.sign_and_send_transaction.return_value = (True, {})
+    staking_contract.extract_events.return_value = [{"name": "ServiceStaked"}]
+    staking_contract.get_staking_state.return_value = StakingState.STAKED
 
-        # We need to make sure prepare_approve_tx is mocked ON THE REGISTRY INSTANCE
-        service_manager.registry.prepare_approve_tx.return_value = {"to": "0xApprove"}
+    # We need to make sure prepare_approve_tx is mocked ON THE REGISTRY INSTANCE
+    service_manager.registry.prepare_approve_tx.return_value = {"to": "0xApprove"}
 
-        assert service_manager.stake(staking_contract) is True
-        assert service_manager.service.staking_contract_address == TEST_STAKING_ADDR
+    assert service_manager.stake(staking_contract) is True
+    assert service_manager.service.staking_contract_address == TEST_STAKING_ADDR
 
 
 def test_unstake_success(service_manager, mock_wallet):
@@ -370,60 +361,15 @@ def test_register_agent_fund_fails(service_manager, mock_wallet):
 
 def test_spin_up_from_pre_registration_success(service_manager, mock_wallet):
     """Test full spin_up path from PRE_REGISTRATION to DEPLOYED."""
-    # Mock state transitions - need to match actual calls in spin_up
-    # The state after activate_registration should be ACTIVE_REGISTRATION
-    state_sequence = [
-        {
-            "state": ServiceState.PRE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # spin_up initial
-        {
-            "state": ServiceState.PRE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # activate_registration check
-        {
-            "state": ServiceState.PRE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # activate_registration internal (get security deposit)
-        {
-            "state": ServiceState.ACTIVE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # spin_up verify after activate
-        {
-            "state": ServiceState.ACTIVE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # register_agent check
-        {
-            "state": ServiceState.ACTIVE_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # register_agent internal
-        {
-            "state": ServiceState.FINISHED_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # spin_up verify after register
-        {
-            "state": ServiceState.FINISHED_REGISTRATION,
-            "security_deposit": 50000000000000000000,
-        },  # deploy check
-        {
-            "state": ServiceState.DEPLOYED,
-            "security_deposit": 50000000000000000000,
-        },  # spin_up verify after deploy
-        {
-            "state": ServiceState.DEPLOYED,
-            "security_deposit": 50000000000000000000,
-        },  # final verification
+    # Use a stateful mock that progresses based on extract_events calls
+    states = [
+        ServiceState.PRE_REGISTRATION,  # Before any TX
+        ServiceState.ACTIVE_REGISTRATION,  # After 1st TX (activate)
+        ServiceState.FINISHED_REGISTRATION,  # After 2nd TX (register)
+        ServiceState.DEPLOYED,  # After 3rd TX (deploy)
     ]
-    service_manager.registry.get_service.side_effect = state_sequence
-
-    mock_wallet.send.return_value = "0xMockTxHash"  # wallet.send returns tx_hash
-    mock_wallet.sign_and_send_transaction.return_value = (True, {})
-
-    # Mock balance/allowance for activate_registration internal call
-    mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
-    mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
-    service_manager.registry.extract_events.side_effect = [
+    tx_count = [0]  # Track completed transactions
+    events_side_effects = [
         [{"name": "ActivateRegistration"}],
         [{"name": "RegisterInstance"}],
         [
@@ -431,6 +377,32 @@ def test_spin_up_from_pre_registration_success(service_manager, mock_wallet):
             {"name": "CreateMultisigWithAgents", "args": {"multisig": TEST_MULTISIG_ADDR}},
         ],
     ]
+    event_idx = [0]
+
+    def dynamic_state(*args, **kwargs):
+        """Return state based on completed transactions."""
+        state = states[min(tx_count[0], len(states) - 1)]
+        return {"state": state, "security_deposit": 50000000000000000000}
+
+    def extract_and_progress(*args, **kwargs):
+        """Return events and advance transaction counter."""
+        if event_idx[0] < len(events_side_effects):
+            events = events_side_effects[event_idx[0]]
+            event_idx[0] += 1
+            tx_count[0] += 1
+            return events
+        return []
+
+    service_manager.registry.get_service.side_effect = dynamic_state
+    service_manager.registry.extract_events.side_effect = extract_and_progress
+
+    mock_wallet.send.return_value = "0xMockTxHash"
+    mock_wallet.sign_and_send_transaction.return_value = (True, {})
+
+    # Mock balance/allowance for activate_registration internal call
+    mock_wallet.balance_service = MagicMock()
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
+    mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     assert service_manager.spin_up() is True
 
@@ -557,14 +529,12 @@ def test_spin_up_with_staking(service_manager, mock_wallet):
         "required_agent_bond": 50000000000000000000,  # 50 OLAS
     }
 
-    with patch("iwa.plugins.olas.service_manager.staking.ERC20Contract") as mock_erc20:
-        mock_erc20.return_value.balance_of_wei.return_value = 100000000000000000000  # 100 OLAS
-        mock_wallet.sign_and_send_transaction.return_value = (True, {})
-        staking_contract.extract_events.return_value = [{"name": "ServiceStaked"}]
-        staking_contract.get_staking_state.return_value = StakingState.STAKED
-        service_manager.registry.prepare_approve_tx.return_value = {"to": "0xApprove"}
+    mock_wallet.sign_and_send_transaction.return_value = (True, {})
+    staking_contract.extract_events.return_value = [{"name": "ServiceStaked"}]
+    staking_contract.get_staking_state.return_value = StakingState.STAKED
+    service_manager.registry.prepare_approve_tx.return_value = {"to": "0xApprove"}
 
-        assert service_manager.spin_up(staking_contract=staking_contract) is True
+    assert service_manager.spin_up(staking_contract=staking_contract) is True
 
 
 def test_spin_up_activate_fails(service_manager, mock_wallet):
@@ -587,7 +557,7 @@ def test_spin_up_activate_fails(service_manager, mock_wallet):
 
     # Mock balance/allowance for activate_registration behavior
     mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     mock_wallet.sign_and_send_transaction.return_value = (False, {})
@@ -931,7 +901,7 @@ def test_activate_registration_token_service_sends_security_deposit_as_value(
 
     # Mock balance check to pass
     mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18  # Plenty of balance
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18  # Plenty of balance
 
     # Mock allowance to pass check (return an int)
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20  # Plenty of allowance
@@ -964,7 +934,7 @@ def test_activate_registration_native_service_sends_security_deposit_as_value(
 
     # Mock balance/allowance
     mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     service_manager.activate_registration()
@@ -996,7 +966,7 @@ def test_activate_registration_uses_master_account_as_from_address(service_manag
     # Mock balance/allowance
     mock_wallet.balance_service = MagicMock()
     mock_wallet.transfer_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     service_manager.activate_registration()
@@ -1028,7 +998,7 @@ def test_activate_registration_uses_master_account_as_signer(service_manager, mo
     # Mock balance/allowance
     mock_wallet.balance_service = MagicMock()
     mock_wallet.transfer_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 10**20
 
     service_manager.activate_registration()
@@ -1055,7 +1025,7 @@ def test_activate_registration_token_service_approves_token_utility(service_mana
 
     # Mock low allowance to trigger approval
     mock_wallet.balance_service = MagicMock()
-    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 10**18
+    mock_wallet.balance_service.get_erc20_balance_wei.return_value = 100 * 10**18
     mock_wallet.transfer_service.get_erc20_allowance.return_value = 0  # Low allowance
     mock_wallet.transfer_service.approve_erc20.return_value = True
 

iwa/plugins/olas/tests/test_service_manager_errors.py

@@ -121,8 +121,7 @@ def test_service_manager_lifecycle_failures(mock_wallet):
         "num_agent_instances": 1,
         "required_agent_bond": 50000000000000000000,
     }
-    with patch("iwa.plugins.olas.service_manager.staking.ERC20Contract"):
-        assert manager.stake(mock_staking) is False
+    assert manager.stake(mock_staking) is False
 
     # unstake failures
     # Service not staked

iwa/plugins/olas/tests/test_service_manager_flows.py

@@ -45,9 +45,8 @@ def mock_config():
 @patch("iwa.plugins.olas.service_manager.base.Config")
 @patch("iwa.plugins.olas.service_manager.base.ServiceRegistryContract")
 @patch("iwa.plugins.olas.service_manager.base.ServiceManagerContract")
-@patch("iwa.plugins.olas.service_manager.staking.ERC20Contract")
 def test_create_service_success(
-    mock_erc20, mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
+    mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
 ):
     """Test successful service creation."""
     # Setup Config with new OlasConfig structure
@@ -85,9 +84,8 @@ def test_create_service_success(
 @patch("iwa.plugins.olas.service_manager.base.Config")
 @patch("iwa.plugins.olas.service_manager.base.ServiceRegistryContract")
 @patch("iwa.plugins.olas.service_manager.base.ServiceManagerContract")
-@patch("iwa.plugins.olas.service_manager.staking.ERC20Contract")
 def test_create_service_failures(
-    mock_erc20, mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
+    mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
 ):
     """Test service creation failure modes."""
     mock_config_inst = mock_config_cls.return_value
@@ -141,9 +139,8 @@ def test_create_service_failures(
 @patch("iwa.plugins.olas.service_manager.base.Config")
 @patch("iwa.plugins.olas.service_manager.base.ServiceRegistryContract")
 @patch("iwa.plugins.olas.service_manager.base.ServiceManagerContract")
-@patch("iwa.plugins.olas.service_manager.staking.ERC20Contract")
 def test_create_service_with_approval(
-    mock_erc20, mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
+    mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet
 ):
     """Test service creation with token approval."""
     mock_config_inst = mock_config_cls.return_value
@@ -404,8 +401,7 @@ def test_terminate(mock_sm_contract, mock_registry_contract, mock_config_cls, mo
 @patch(
     "iwa.plugins.olas.service_manager.base.ServiceManagerContract"
 )  # MUST mock specifically here
-@patch("iwa.plugins.olas.service_manager.staking.ERC20Contract")  # For checking balance
-def test_stake(mock_erc20, mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet):
+def test_stake(mock_sm_contract, mock_registry_contract, mock_config_cls, mock_wallet):
     """Test service staking."""
     # Setup mock service
     mock_service = MagicMock()
@@ -445,9 +441,6 @@ def test_stake(mock_erc20, mock_sm_contract, mock_registry_contract, mock_config
         "required_agent_bond": 50000000000000000000,
     }
 
-    # Mock ERC20 balance check (100 OLAS = 100e18 wei, enough for 50 min deposit)
-    mock_erc20_inst = mock_erc20.return_value
-    mock_erc20_inst.balance_of_wei.return_value = 100000000000000000000  # 100 OLAS
 
     success = manager.stake(mock_staking)
     assert success is True
@@ -471,10 +464,7 @@ def test_stake(mock_erc20, mock_sm_contract, mock_registry_contract, mock_config
     assert manager.stake(mock_staking) is False
     mock_staking.get_service_ids.return_value = []
 
-    # 3. Not enough funds
-    mock_erc20_inst.balance_of_wei.return_value = 50
-    assert manager.stake(mock_staking) is False
-    mock_erc20_inst.balance_of_wei.return_value = 200
+
 
     # 4. Approve fail
     mock_wallet.sign_and_send_transaction.return_value = (False, {})