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

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

tests/test_safe_executor.py

@@ -567,3 +567,73 @@ def test_rpc_rotation_stops_when_should_not_retry(
         assert success is False
         # Only 1 attempt because should_retry=False
         assert mock_safe_tx.execute.call_count == 1
+
+
+# =============================================================================
+# Test: Fee bumping on base fee errors
+# =============================================================================
+
+
+def test_fee_error_triggers_bump(executor, mock_chain_interface, mock_safe_tx, mock_safe):
+    """Test that fee errors trigger gas price bump on retry."""
+    with patch.object(executor, "_recreate_safe_client", return_value=mock_safe):
+        # First attempt fails with fee error, second succeeds
+        mock_safe_tx.execute.side_effect = [
+            ValueError("max fee per gas less than block base fee: maxFeePerGas: 596, baseFee: 681"),
+            b"tx_hash",
+        ]
+        mock_chain_interface.web3.eth.wait_for_transaction_receipt.return_value = MagicMock(
+            status=1
+        )
+        # Mock fee calculation
+        mock_chain_interface.web3.eth.get_block.return_value = {"baseFeePerGas": 700}
+        mock_chain_interface.web3.eth.max_priority_fee = 1
+
+        with patch("time.sleep"):
+            success, tx_hash, receipt = executor.execute_with_retry(
+                "0xSafe", mock_safe_tx, ["key1"]
+            )
+
+        assert success is True
+        assert mock_safe_tx.execute.call_count == 2
+        # Second call should have tx_gas_price (bumped), not eip1559_speed
+        second_call_kwargs = mock_safe_tx.execute.call_args_list[1][1]
+        assert "tx_gas_price" in second_call_kwargs
+
+
+def test_fee_error_classification(executor):
+    """Test classification of fee-related errors."""
+    fee_errors = [
+        "max fee per gas less than block base fee",
+        "transaction underpriced",
+        "maxFeePerGas too low",
+        "fee too low for mempool",
+    ]
+    for error_msg in fee_errors:
+        error = ValueError(error_msg)
+        result = executor._classify_error(error)
+        assert result["is_fee_error"] is True, f"Should detect fee error: {error_msg}"
+
+
+def test_calculate_bumped_gas_price_eip1559(executor, mock_chain_interface):
+    """Test bumped gas price calculation for EIP-1559 chains."""
+    mock_chain_interface.web3.eth.get_block.return_value = {"baseFeePerGas": 1000}
+    mock_chain_interface.web3.eth.max_priority_fee = 10
+
+    # With 1.3x bump factor: base_fee * 1.3 * 1.5 + priority = 1000 * 1.3 * 1.5 + 10 = 1960
+    result = executor._calculate_bumped_gas_price(1.3)
+
+    assert result is not None
+    assert result == int(1000 * 1.3 * 1.5) + 10
+
+
+def test_calculate_bumped_gas_price_legacy(executor, mock_chain_interface):
+    """Test bumped gas price calculation for legacy chains."""
+    mock_chain_interface.web3.eth.get_block.return_value = {}  # No baseFeePerGas
+    mock_chain_interface.web3.eth.gas_price = 2000
+
+    # Legacy: gas_price * bump_factor = 2000 * 1.3 = 2600
+    result = executor._calculate_bumped_gas_price(1.3)
+
+    assert result is not None
+    assert result == int(2000 * 1.3)