iwa 0.0.0__py3-none-any.whl → 0.0.1a2__py3-none-any.whl

iwa/plugins/olas/tests/test_service_manager_rewards.py

@@ -0,0 +1,360 @@
+"""Tests for staking rewards claim, withdraw, and checkpoint functionality."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from iwa.plugins.olas.contracts.staking import StakingState
+from iwa.plugins.olas.models import Service
+from iwa.plugins.olas.service_manager import ServiceManager
+
+VALID_ADDR = "0x1234567890123456789012345678901234567890"
+WITHDRAWAL_ADDR = "0xABCDEFabcdefABCDEFabcdefABCDEFabcdefABCD"
+
+
+@pytest.fixture
+def mock_wallet():
+    """Create a mock wallet for tests."""
+    wallet = MagicMock()
+    wallet.master_account.address = VALID_ADDR
+    wallet.chain_interface.chain_name = "gnosis"
+    wallet.sign_and_send_transaction.return_value = (True, {"transactionHash": b"\x00" * 32})
+    wallet.send.return_value = "0xtxhash"
+    return wallet
+
+
+def setup_manager(mock_wallet):
+    """Setup a ServiceManager with mocked dependencies."""
+    with patch("iwa.plugins.olas.service_manager.base.Config") as mock_cfg_cls:
+        mock_cfg = mock_cfg_cls.return_value
+        mock_olas_config = MagicMock()
+        mock_olas_config.get_service.return_value = None
+        mock_olas_config.withdrawal_address = None
+        mock_cfg.plugins = {"olas": mock_olas_config}
+        with patch(
+            "iwa.plugins.olas.service_manager.OLAS_CONTRACTS",
+            {
+                "gnosis": {
+                    "OLAS_SERVICE_REGISTRY": VALID_ADDR,
+                    "OLAS_SERVICE_MANAGER": VALID_ADDR,
+                }
+            },
+        ):
+            with patch("iwa.plugins.olas.service_manager.base.ChainInterfaces") as mock_if_cls:
+                mock_if = mock_if_cls.return_value
+                mock_if.get.return_value.chain.name.lower.return_value = "gnosis"
+                mock_if.get.return_value.get_contract_address.return_value = VALID_ADDR
+                manager = ServiceManager(mock_wallet)
+                manager.registry = MagicMock()
+                manager.manager_contract = MagicMock()
+                manager.olas_config = mock_olas_config
+                manager.chain_name = "gnosis"
+                return manager
+
+
+def test_claim_rewards_no_service(mock_wallet):
+    """Test claim_rewards fails when no service is loaded."""
+    manager = setup_manager(mock_wallet)
+    manager.service = None
+
+    success, amount = manager.claim_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_claim_rewards_not_staked(mock_wallet):
+    """Test claim_rewards fails when service is not staked."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(service_name="test", chain_name="gnosis", service_id=1, agent_ids=[1])
+    manager.service.staking_contract_address = None
+
+    success, amount = manager.claim_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_claim_rewards_service_not_staked_state(mock_wallet):
+    """Test claim_rewards fails when staking state is NOT_STAKED."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.get_staking_state.return_value = StakingState.NOT_STAKED
+
+    success, amount = manager.claim_rewards(staking_contract=mock_staking)
+
+    assert success is False
+    assert amount == 0
+
+
+def test_claim_rewards_no_accrued_rewards(mock_wallet):
+    """Test claim_rewards fails when no rewards are accrued."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.get_staking_state.return_value = StakingState.STAKED
+    mock_staking.get_accrued_rewards.return_value = 0
+
+    success, amount = manager.claim_rewards(staking_contract=mock_staking)
+
+    assert success is False
+    assert amount == 0
+
+
+def test_claim_rewards_success(mock_wallet):
+    """Test claim_rewards succeeds with proper setup."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.get_staking_state.return_value = StakingState.STAKED
+    mock_staking.get_accrued_rewards.return_value = 10 * 10**18  # 10 OLAS
+    mock_staking.prepare_claim_tx.return_value = {"data": "0x"}
+    mock_staking.extract_events.return_value = [{"name": "RewardClaimed"}]
+
+    success, amount = manager.claim_rewards(staking_contract=mock_staking)
+
+    assert success is True
+    assert amount == 10 * 10**18
+
+
+def test_claim_rewards_tx_fails(mock_wallet):
+    """Test claim_rewards fails when transaction fails."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.get_staking_state.return_value = StakingState.STAKED
+    mock_staking.get_accrued_rewards.return_value = 10 * 10**18
+    mock_staking.prepare_claim_tx.return_value = {"data": "0x"}
+    mock_wallet.sign_and_send_transaction.return_value = (False, {})
+
+    success, amount = manager.claim_rewards(staking_contract=mock_staking)
+
+    assert success is False
+    assert amount == 0
+
+
+def test_withdraw_rewards_no_service(mock_wallet):
+    """Test withdraw_rewards fails when no service is loaded."""
+    manager = setup_manager(mock_wallet)
+    manager.service = None
+
+    success, amount = manager.withdraw_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_withdraw_rewards_no_multisig(mock_wallet):
+    """Test withdraw_rewards fails when service has no multisig."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(service_name="test", chain_name="gnosis", service_id=1, agent_ids=[1])
+    manager.service.multisig_address = None
+
+    success, amount = manager.withdraw_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_withdraw_rewards_no_withdrawal_address(mock_wallet):
+    """Test withdraw_rewards fails when no withdrawal address configured."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        multisig_address=VALID_ADDR,
+    )
+    manager.olas_config.withdrawal_address = None
+
+    success, amount = manager.withdraw_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_withdraw_rewards_no_olas_balance(mock_wallet):
+    """Test withdraw_rewards fails when Safe has no OLAS balance."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        multisig_address=VALID_ADDR,
+    )
+    manager.olas_config.withdrawal_address = WITHDRAWAL_ADDR
+
+    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 = 0
+
+        success, amount = manager.withdraw_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+def test_withdraw_rewards_success(mock_wallet):
+    """Test withdraw_rewards succeeds with proper setup."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        multisig_address=VALID_ADDR,
+    )
+    manager.olas_config.withdrawal_address = WITHDRAWAL_ADDR
+
+    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 = 50 * 10**18  # 50 OLAS
+
+        success, amount = manager.withdraw_rewards()
+
+    assert success is True
+    assert amount == 50.0
+    mock_wallet.send.assert_called_once()
+
+
+def test_withdraw_rewards_transfer_fails(mock_wallet):
+    """Test withdraw_rewards fails when transfer fails."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        multisig_address=VALID_ADDR,
+    )
+    manager.olas_config.withdrawal_address = WITHDRAWAL_ADDR
+    mock_wallet.send.return_value = None  # Transfer fails
+
+    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 = 50 * 10**18
+
+        success, amount = manager.withdraw_rewards()
+
+    assert success is False
+    assert amount == 0
+
+
+# ============================================================================
+# Checkpoint Tests
+# ============================================================================
+
+
+def test_call_checkpoint_no_service(mock_wallet):
+    """Test call_checkpoint fails when no service is loaded."""
+    manager = setup_manager(mock_wallet)
+    manager.service = None
+
+    result = manager.call_checkpoint()
+
+    assert result is False
+
+
+def test_call_checkpoint_not_staked(mock_wallet):
+    """Test call_checkpoint fails when service is not staked."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(service_name="test", chain_name="gnosis", service_id=1, agent_ids=[1])
+    manager.service.staking_contract_address = None
+
+    result = manager.call_checkpoint()
+
+    assert result is False
+
+
+def test_call_checkpoint_not_needed(mock_wallet):
+    """Test call_checkpoint returns False when checkpoint is not needed."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.is_checkpoint_needed.return_value = False
+    mock_staking.get_next_epoch_start.return_value = MagicMock()
+
+    result = manager.call_checkpoint(staking_contract=mock_staking)
+
+    assert result is False
+
+
+def test_call_checkpoint_success(mock_wallet):
+    """Test call_checkpoint succeeds with proper setup."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.is_checkpoint_needed.return_value = True
+    mock_staking.prepare_checkpoint_tx.return_value = {"data": "0x", "gas": 4_000_000}
+    mock_staking.extract_events.return_value = [{"name": "Checkpoint"}]
+
+    result = manager.call_checkpoint(staking_contract=mock_staking)
+
+    assert result is True
+    mock_staking.prepare_checkpoint_tx.assert_called_once()
+
+
+def test_call_checkpoint_tx_fails(mock_wallet):
+    """Test call_checkpoint fails when transaction fails."""
+    manager = setup_manager(mock_wallet)
+    manager.service = Service(
+        service_name="test",
+        chain_name="gnosis",
+        service_id=1,
+        agent_ids=[1],
+        staking_contract_address=VALID_ADDR,
+    )
+
+    mock_staking = MagicMock()
+    mock_staking.is_checkpoint_needed.return_value = True
+    mock_staking.prepare_checkpoint_tx.return_value = {"data": "0x", "gas": 4_000_000}
+    mock_wallet.sign_and_send_transaction.return_value = (False, {})
+
+    result = manager.call_checkpoint(staking_contract=mock_staking)
+
+    assert result is False

iwa/plugins/olas/tests/test_service_manager_validation.py

@@ -0,0 +1,145 @@
+"""Tests for Olas ServiceManager validation and failure handling."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from iwa.plugins.olas.contracts.service import ServiceState
+from iwa.plugins.olas.models import Service
+from iwa.plugins.olas.service_manager import ServiceManager
+
+VALID_ADDR_1 = "0x78731D3Ca6b7E34aC0F824c42a7cC18A495cabaB"
+VALID_ADDR_2 = "0x40A2aCCbd92BCA938b02010E17A5b8929b49130D"
+
+
+@pytest.fixture
+def mock_wallet():
+    """Mock wallet."""
+    wallet = MagicMock()
+    wallet.master_account.address = VALID_ADDR_1
+    wallet.transfer_service = MagicMock()
+    return wallet
+
+
+@pytest.fixture
+def sm(mock_wallet):
+    """ServiceManager fixture."""
+    with patch("iwa.core.models.Config"):
+        manager = ServiceManager(mock_wallet)
+        # Mock service
+        manager.service = Service(
+            service_id=1,
+            service_name="Test",
+            chain_name="gnosis",
+            agent_address=VALID_ADDR_1,
+            multisig_address=VALID_ADDR_2,
+            staking_contract_address=VALID_ADDR_1,
+        )
+        return manager
+
+
+def test_drain_service_partial_failures(sm, mock_wallet):
+    """Test drain_service handles partial failures across accounts."""
+    # Setup:
+    # 1. Claim success
+    # 2. Safe drain failure
+    # 3. Agent drain success
+
+    with patch.object(sm, "claim_rewards", return_value=(True, 10**18)):
+        # Wallet.drain is called for Safe and Agent
+        def mock_drain(from_address_or_tag=None, to_address_or_tag=None, chain_name=None):
+            if from_address_or_tag == VALID_ADDR_2:  # Safe
+                raise Exception("Safe drain failed")
+            return {"native": 0.5}
+
+        mock_wallet.drain.side_effect = mock_drain
+
+        result = sm.drain_service()
+
+        assert "safe" not in result
+        assert "agent" in result
+        assert result["agent"]["native"] == 0.5
+        # Verify it continued after Safe failure
+
+
+def test_unstake_failed_event_extraction(sm):
+    """Test unstake when transaction succeeds but event extraction fails."""
+    staking_mock = MagicMock()
+    sm.wallet.sign_and_send_transaction.return_value = (True, {"tx_hash": "0x123"})
+    staking_mock.extract_events.return_value = []  # No Unstaked event
+
+    success = sm.unstake(staking_mock)
+    assert success is False  # Should return False if event missing
+
+
+def test_call_checkpoint_grace_period(sm):
+    """Test call_checkpoint respect for grace period."""
+    staking_mock = MagicMock()
+    # Mock status where epoch ended very recently (within grace period)
+    staking_mock.get_staking_status.return_value = {
+        "remaining_epoch_seconds": -50  # Ended 50s ago
+    }
+    # Mock is_checkpoint_needed to return False based on grace period
+    staking_mock.is_checkpoint_needed.return_value = False
+    staking_mock.get_next_epoch_start.return_value = MagicMock()
+
+    # grace_period_seconds defaults to 600
+    success = sm.call_checkpoint(staking_mock, grace_period_seconds=600)
+
+    # Should skip checkpoint and return False (matching SM logic)
+    assert success is False
+    sm.wallet.sign_and_send_transaction.assert_not_called()
+
+
+def test_call_checkpoint_success(sm):
+    """Test successful call_checkpoint."""
+    staking_mock = MagicMock()
+    staking_mock.is_checkpoint_needed.return_value = True
+    staking_mock.prepare_checkpoint_tx.return_value = {"to": "0x123", "data": "0x"}
+    sm.wallet.sign_and_send_transaction.return_value = (True, {"tx_hash": "0x123", "logs": []})
+    staking_mock.extract_events.return_value = [
+        {"name": "Checkpoint", "args": {"epoch": 1, "availableRewards": 10**18}}
+    ]
+
+    success = sm.call_checkpoint(staking_mock)
+    assert success is True
+    sm.wallet.sign_and_send_transaction.assert_called_once()
+
+
+def test_spin_up_intermediate_failure(sm):
+    """Test spin_up stops at first failure."""
+    # Mock sequential calls using ServiceState enum on the registry mock
+    with patch.object(
+        sm.registry,
+        "get_service",
+        side_effect=[
+            {"state": ServiceState.PRE_REGISTRATION},
+            {"state": ServiceState.ACTIVE_REGISTRATION},
+            {"state": ServiceState.ACTIVE_REGISTRATION},  # Verification after activate
+            {"state": ServiceState.ACTIVE_REGISTRATION},  # Final verification (if it reached there)
+            {"state": ServiceState.ACTIVE_REGISTRATION},  # One more just in case
+        ],
+    ):
+        with (
+            patch.object(sm, "activate_registration", return_value=True) as m1,
+            patch.object(sm, "register_agent", return_value=False) as m2,
+            patch.object(sm, "deploy") as m3,
+        ):
+            success = sm.spin_up()
+
+            assert success is False
+            m1.assert_called_once()
+            m2.assert_called_once()
+            m3.assert_not_called()
+
+
+def test_service_manager_no_service_error_handling(mock_wallet):
+    """Test methods return gracefully when no service is selected."""
+    with patch("iwa.core.models.Config"):
+        manager = ServiceManager(mock_wallet)
+        manager.service = None
+
+        assert manager.get() is None
+        assert manager.drain_service() == {}
+        assert manager.get_staking_status() is None
+        assert manager.claim_rewards() == (False, 0)