iwa 0.0.62__py3-none-any.whl → 0.0.64__py3-none-any.whl

tests/test_contract_cache.py

@@ -0,0 +1,253 @@
+"""Tests for contract instance caching."""
+
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def reset_singleton():
+    """Reset the ContractCache singleton before each test."""
+    from iwa.core.contracts.cache import ContractCache
+
+    ContractCache._instance = None
+    yield
+    ContractCache._instance = None
+
+
+class TestContractCache:
+    """Test ContractCache singleton and caching behavior."""
+
+    def test_singleton_pattern(self):
+        """Test that ContractCache is a singleton."""
+        from iwa.core.contracts.cache import ContractCache
+
+        c1 = ContractCache()
+        c2 = ContractCache()
+        assert c1 is c2
+
+    def test_get_contract_creates_new_instance(self):
+        """Test get_contract creates new instance when not cached."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_instance = MagicMock()
+        mock_cls.return_value = mock_instance
+
+        result = cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        mock_cls.assert_called_once_with(
+            "0x1234567890123456789012345678901234567890", chain_name="gnosis"
+        )
+        assert result is mock_instance
+
+    def test_get_contract_returns_cached_instance(self):
+        """Test get_contract returns cached instance on second call."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_instance = MagicMock()
+        mock_cls.return_value = mock_instance
+
+        result1 = cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+        result2 = cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        # Should only create once
+        mock_cls.assert_called_once()
+        assert result1 is result2
+
+    def test_get_contract_raises_on_empty_address(self):
+        """Test get_contract raises ValueError for empty address."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+
+        with pytest.raises(ValueError, match="Address is required"):
+            cache.get_contract(mock_cls, "", "gnosis")
+
+    def test_get_contract_respects_ttl_expiry(self):
+        """Test get_contract recreates instance after TTL expires."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_cls.return_value = MagicMock()
+
+        # First call
+        cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis", ttl=0
+        )
+
+        # Wait for expiry (TTL=0 means immediate expiry)
+        time.sleep(0.01)
+
+        # Second call should create new instance
+        cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis", ttl=0
+        )
+
+        assert mock_cls.call_count == 2
+
+    def test_get_if_cached_returns_cached_instance(self):
+        """Test get_if_cached returns cached instance."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_instance = MagicMock()
+        mock_cls.return_value = mock_instance
+
+        # First populate cache
+        cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        # get_if_cached should return it
+        result = cache.get_if_cached(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        assert result is mock_instance
+
+    def test_get_if_cached_returns_none_when_not_cached(self):
+        """Test get_if_cached returns None when not cached."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+
+        result = cache.get_if_cached(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        assert result is None
+
+    def test_get_if_cached_returns_none_for_empty_address(self):
+        """Test get_if_cached returns None for empty address."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+
+        result = cache.get_if_cached(mock_cls, "", "gnosis")
+
+        assert result is None
+
+    def test_get_if_cached_returns_none_after_expiry(self):
+        """Test get_if_cached returns None after TTL expires."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        cache.ttl = 0  # Immediate expiry
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_cls.return_value = MagicMock()
+
+        # Populate cache
+        cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        time.sleep(0.01)
+
+        # get_if_cached should return None due to expiry
+        result = cache.get_if_cached(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        assert result is None
+
+    def test_clear_removes_all_entries(self):
+        """Test clear removes all cached contracts."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="MockContract")
+        mock_cls.return_value = MagicMock()
+
+        # Populate cache
+        cache.get_contract(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        cache.clear()
+
+        # get_if_cached should return None
+        result = cache.get_if_cached(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+        assert result is None
+
+    def test_invalidate_removes_specific_entry(self):
+        """Test invalidate removes specific cached contract."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls1 = MagicMock(__name__="Contract1")
+        mock_cls2 = MagicMock(__name__="Contract2")
+        mock_cls1.return_value = MagicMock()
+        mock_cls2.return_value = MagicMock()
+
+        # Populate cache with two contracts
+        cache.get_contract(
+            mock_cls1, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+        cache.get_contract(
+            mock_cls2, "0xABCDEF1234567890123456789012345678901234", "gnosis"
+        )
+
+        # Invalidate only the first one
+        cache.invalidate(
+            mock_cls1, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+        # First should be gone
+        result1 = cache.get_if_cached(
+            mock_cls1, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+        assert result1 is None
+
+        # Second should still exist
+        result2 = cache.get_if_cached(
+            mock_cls2, "0xABCDEF1234567890123456789012345678901234", "gnosis"
+        )
+        assert result2 is not None
+
+    def test_invalidate_nonexistent_does_nothing(self):
+        """Test invalidate does nothing for non-existent entry."""
+        from iwa.core.contracts.cache import ContractCache
+
+        cache = ContractCache()
+        mock_cls = MagicMock(__name__="Contract")
+
+        # Should not raise
+        cache.invalidate(
+            mock_cls, "0x1234567890123456789012345678901234567890", "gnosis"
+        )
+
+    def test_env_ttl_configuration(self):
+        """Test TTL is configurable via environment variable."""
+        from iwa.core.contracts.cache import ContractCache
+
+        with patch.dict("os.environ", {"IWA_CONTRACT_CACHE_TTL": "7200"}):
+            ContractCache._instance = None
+            cache = ContractCache()
+            assert cache.ttl == 7200
+
+    def test_invalid_env_ttl_uses_default(self):
+        """Test invalid TTL env var uses default value."""
+        from iwa.core.contracts.cache import ContractCache
+
+        with patch.dict("os.environ", {"IWA_CONTRACT_CACHE_TTL": "invalid"}):
+            ContractCache._instance = None
+            cache = ContractCache()
+            assert cache.ttl == 3600

tests/test_drain_coverage.py

@@ -46,7 +46,7 @@ def test_claim_rewards_claim_tx_fails(mock_drain_manager):
     with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
         mock_staking = mock_staking_cls.return_value
         mock_staking.get_staking_state.return_value = StakingState.STAKED
-        mock_staking.get_accrued_rewards.return_value = 1000000000000000000
+        mock_staking.calculate_staking_reward.return_value = 1000000000000000000
         mock_staking.prepare_claim_tx.return_value = None  # Failed to prepare
 
         success, amount = mock_drain_manager.claim_rewards()
@@ -63,7 +63,7 @@ def test_claim_rewards_send_fails(mock_drain_manager):
     with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
         mock_staking = mock_staking_cls.return_value
         mock_staking.get_staking_state.return_value = StakingState.STAKED
-        mock_staking.get_accrued_rewards.return_value = 1000000000000000000
+        mock_staking.calculate_staking_reward.return_value = 1000000000000000000
         mock_staking.prepare_claim_tx.return_value = {"to": "0x", "data": "0x"}
         mock_drain_manager.wallet.sign_and_send_transaction.return_value = (False, None)
 
@@ -81,7 +81,7 @@ def test_claim_rewards_success_no_event(mock_drain_manager):
     with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
         mock_staking = mock_staking_cls.return_value
         mock_staking.get_staking_state.return_value = StakingState.STAKED
-        mock_staking.get_accrued_rewards.return_value = 1000000000000000000
+        mock_staking.calculate_staking_reward.return_value = 1000000000000000000
         mock_staking.prepare_claim_tx.return_value = {"to": "0x", "data": "0x"}
         mock_staking.extract_events.return_value = []  # No RewardClaimed event
         mock_drain_manager.wallet.sign_and_send_transaction.return_value = (
@@ -94,6 +94,33 @@ def test_claim_rewards_success_no_event(mock_drain_manager):
         assert amount == 1000000000000000000
 
 
+def test_claim_rewards_fallback_to_accrued(mock_drain_manager):
+    """Test claim_rewards falls back to get_accrued_rewards when calculate_staking_reward fails."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+    mock_drain_manager.wallet.master_account.address = "0xMaster"
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking = mock_staking_cls.return_value
+        mock_staking.get_staking_state.return_value = StakingState.STAKED
+        # calculate_staking_reward fails, should fallback to get_accrued_rewards
+        mock_staking.calculate_staking_reward.side_effect = Exception("RPC error")
+        mock_staking.get_accrued_rewards.return_value = 2000000000000000000  # 2 OLAS
+        mock_staking.prepare_claim_tx.return_value = {"to": "0x", "data": "0x"}
+        mock_staking.extract_events.return_value = []
+        mock_drain_manager.wallet.sign_and_send_transaction.return_value = (
+            True,
+            {"transactionHash": "0xHash"},
+        )
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert success
+        assert amount == 2000000000000000000
+        # Verify both methods were called
+        mock_staking.calculate_staking_reward.assert_called_once()
+        mock_staking.get_accrued_rewards.assert_called_once()
+
+
 def test_withdraw_rewards_fallback_to_master(mock_drain_manager):
     """Test withdraw_rewards falls back to master account when not configured."""
     mock_drain_manager.service.multisig_address = "0x1111111111111111111111111111111111111111"
@@ -178,3 +205,238 @@ def test_drain_owner_skipped_when_equals_target(mock_drain_manager):
     # Should skip and return None without calling drain
     assert result is None
     mock_drain_manager.wallet.drain.assert_not_called()
+
+
+def test_claim_rewards_staking_contract_load_fails(mock_drain_manager):
+    """Test claim_rewards when StakingContract fails to load."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking_cls.side_effect = Exception("Failed to load contract")
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert not success
+        assert amount == 0
+
+
+def test_claim_rewards_not_staked_state(mock_drain_manager):
+    """Test claim_rewards when service staking state is not STAKED."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking = mock_staking_cls.return_value
+        mock_staking.get_staking_state.return_value = StakingState.NOT_STAKED
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert not success
+        assert amount == 0
+
+
+def test_claim_rewards_zero_rewards(mock_drain_manager):
+    """Test claim_rewards when claimable rewards is zero."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking = mock_staking_cls.return_value
+        mock_staking.get_staking_state.return_value = StakingState.STAKED
+        mock_staking.calculate_staking_reward.return_value = 0
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert not success
+        assert amount == 0
+
+
+def test_claim_rewards_with_event_amount(mock_drain_manager):
+    """Test claim_rewards extracts amount from RewardClaimed event."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+    mock_drain_manager.wallet.master_account.address = "0xMaster"
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking = mock_staking_cls.return_value
+        mock_staking.get_staking_state.return_value = StakingState.STAKED
+        mock_staking.calculate_staking_reward.return_value = 1000000000000000000
+        mock_staking.prepare_claim_tx.return_value = {"to": "0x", "data": "0x"}
+        # RewardClaimed event with 'amount' field
+        mock_staking.extract_events.return_value = [
+            {"name": "RewardClaimed", "args": {"amount": 1500000000000000000}}
+        ]
+        mock_drain_manager.wallet.sign_and_send_transaction.return_value = (
+            True,
+            {"transactionHash": "0xHash"},
+        )
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert success
+        # Should use amount from event, not the estimate
+        assert amount == 1500000000000000000
+
+
+def test_claim_rewards_with_event_reward_field(mock_drain_manager):
+    """Test claim_rewards extracts from 'reward' field when 'amount' missing."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.service.service_id = 1
+    mock_drain_manager.wallet.master_account.address = "0xMaster"
+
+    with patch("iwa.plugins.olas.service_manager.drain.StakingContract") as mock_staking_cls:
+        mock_staking = mock_staking_cls.return_value
+        mock_staking.get_staking_state.return_value = StakingState.STAKED
+        mock_staking.calculate_staking_reward.return_value = 1000000000000000000
+        mock_staking.prepare_claim_tx.return_value = {"to": "0x", "data": "0x"}
+        # RewardClaimed event with 'reward' field (no 'amount')
+        mock_staking.extract_events.return_value = [
+            {"name": "RewardClaimed", "args": {"reward": 2000000000000000000}}
+        ]
+        mock_drain_manager.wallet.sign_and_send_transaction.return_value = (
+            True,
+            {"transactionHash": "0xHash"},
+        )
+
+        success, amount = mock_drain_manager.claim_rewards()
+        assert success
+        assert amount == 2000000000000000000
+
+
+def test_withdraw_rewards_no_service(mock_drain_manager):
+    """Test withdraw_rewards with no active service."""
+    mock_drain_manager.service = None
+    success, amount = mock_drain_manager.withdraw_rewards()
+    assert not success
+    assert amount == 0
+
+
+def test_withdraw_rewards_no_multisig(mock_drain_manager):
+    """Test withdraw_rewards when service has no multisig."""
+    mock_drain_manager.service.multisig_address = None
+    success, amount = mock_drain_manager.withdraw_rewards()
+    assert not success
+    assert amount == 0
+
+
+def test_withdraw_rewards_success(mock_drain_manager):
+    """Test withdraw_rewards succeeds with balance."""
+    mock_drain_manager.service.multisig_address = "0xSafe"
+    mock_drain_manager.olas_config.withdrawal_address = "0xWithdraw"
+    mock_drain_manager.wallet.account_service.get_tag_by_address.return_value = None
+    mock_drain_manager.wallet.send.return_value = "0xTxHash"
+
+    with patch("iwa.plugins.olas.service_manager.drain.ERC20Contract") as mock_erc20_cls:
+        mock_erc20 = mock_erc20_cls.return_value
+        mock_erc20.balance_of_wei.return_value = 5000000000000000000  # 5 OLAS
+
+        success, amount = mock_drain_manager.withdraw_rewards()
+        assert success
+        assert amount == 5.0  # 5 OLAS
+
+
+def test_withdraw_rewards_transfer_fails(mock_drain_manager):
+    """Test withdraw_rewards when transfer fails."""
+    mock_drain_manager.service.multisig_address = "0xSafe"
+    mock_drain_manager.olas_config.withdrawal_address = "0xWithdraw"
+    mock_drain_manager.wallet.account_service.get_tag_by_address.return_value = None
+    mock_drain_manager.wallet.send.return_value = None  # Transfer failed
+
+    with patch("iwa.plugins.olas.service_manager.drain.ERC20Contract") as mock_erc20_cls:
+        mock_erc20 = mock_erc20_cls.return_value
+        mock_erc20.balance_of_wei.return_value = 5000000000000000000
+
+        success, amount = mock_drain_manager.withdraw_rewards()
+        assert not success
+        assert amount == 0
+
+
+def test_drain_service_full_flow(mock_drain_manager):
+    """Test drain_service drains all accounts."""
+    mock_drain_manager.service.key = "test_service"
+    mock_drain_manager.service.staking_contract_address = None  # No staking
+    mock_drain_manager.service.multisig_address = "0xSafe"
+    mock_drain_manager.service.agent_address = "0xAgent"
+    mock_drain_manager.service.service_owner_address = "0xOwner"
+    mock_drain_manager.wallet.master_account.address = "0xMaster"
+
+    # Mock drain returns
+    mock_drain_manager.wallet.drain.return_value = (True, {"transactionHash": b"0x123"})
+
+    result = mock_drain_manager.drain_service(target_address="0xTarget")
+
+    assert "safe" in result or "agent" in result or "owner" in result
+
+
+def test_drain_safe_account_with_retry(mock_drain_manager):
+    """Test _drain_safe_account retries when rewards were claimed."""
+    mock_drain_manager.service.multisig_address = "0xSafe"
+    # First call returns None (balance not updated yet), second returns result
+    mock_drain_manager.wallet.drain.side_effect = [
+        None,
+        (True, {"transactionHash": b"0xhash"}),
+    ]
+
+    with patch("time.sleep"):  # Don't actually sleep
+        mock_drain_manager._drain_safe_account(
+            "0xTarget", "gnosis", claimed_rewards=1000
+        )
+
+    # Should have retried
+    assert mock_drain_manager.wallet.drain.call_count >= 1
+
+
+def test_drain_agent_account_success(mock_drain_manager):
+    """Test _drain_agent_account success."""
+    mock_drain_manager.service.agent_address = "0xAgent"
+    mock_drain_manager.wallet.drain.return_value = (True, {"transactionHash": b"0xhash"})
+
+    result = mock_drain_manager._drain_agent_account("0xTarget", "gnosis")
+
+    assert result is not None
+
+
+def test_drain_owner_account_success(mock_drain_manager):
+    """Test _drain_owner_account success."""
+    mock_drain_manager.service.service_owner_address = "0xOwner"
+    mock_drain_manager.wallet.drain.return_value = (True, {"transactionHash": b"0xhash"})
+
+    result = mock_drain_manager._drain_owner_account("0xTarget", "gnosis")
+
+    assert result is not None
+
+
+def test_normalize_drain_result_passthrough(mock_drain_manager):
+    """Test _normalize_drain_result passes through non-tuple results."""
+    result = mock_drain_manager._normalize_drain_result({"some": "dict"})
+    assert result == {"some": "dict"}
+
+
+def test_normalize_drain_result_string_hash(mock_drain_manager):
+    """Test _normalize_drain_result handles string transaction hash."""
+    result = mock_drain_manager._normalize_drain_result(
+        (True, {"transactionHash": "0xABCDEF"})
+    )
+    assert result == "0xABCDEF"
+
+
+def test_claim_rewards_if_needed_success(mock_drain_manager):
+    """Test _claim_rewards_if_needed returns claimed amount on success."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+    mock_drain_manager.claim_rewards = MagicMock(return_value=(True, 1000000000000000000))
+
+    result = mock_drain_manager._claim_rewards_if_needed(claim_rewards=True)
+    assert result == 1000000000000000000
+
+
+def test_claim_rewards_if_needed_disabled(mock_drain_manager):
+    """Test _claim_rewards_if_needed returns 0 when disabled."""
+    mock_drain_manager.service.staking_contract_address = "0xStaking"
+
+    result = mock_drain_manager._claim_rewards_if_needed(claim_rewards=False)
+    assert result == 0
+
+
+def test_claim_rewards_if_needed_no_staking(mock_drain_manager):
+    """Test _claim_rewards_if_needed returns 0 when not staked."""
+    mock_drain_manager.service.staking_contract_address = None
+
+    result = mock_drain_manager._claim_rewards_if_needed(claim_rewards=True)
+    assert result == 0