wayfinder-paths 0.1.19__py3-none-any.whl → 0.1.21__py3-none-any.whl

wayfinder_paths/strategies/moonwell_wsteth_loop_strategy/test_strategy.py

@@ -1,5 +1,3 @@
-"""Test template for Moonwell wstETH Loop Strategy."""
-
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -25,7 +23,6 @@ from wayfinder_paths.tests.test_utils import (
 
 @pytest.fixture
 def strategy():
-    """Create a strategy instance for testing with minimal config."""
     mock_config = {
         "main_wallet": {"address": "0x1234567890123456789012345678901234567890"},
         "strategy_wallet": {"address": "0xabcdefabcdefabcdefabcdefabcdefabcdefabcd"},
@@ -39,11 +36,9 @@ def strategy():
             config=mock_config,
             main_wallet=mock_config["main_wallet"],
             strategy_wallet=mock_config["strategy_wallet"],
-            simulation=True,
         )
         # Manually set attributes that would be set in __init__
         s.config = mock_config
-        s.simulation = True
         s._token_info_cache = {}
         s._token_price_cache = {}
         s._token_price_timestamps = {}
@@ -61,7 +56,6 @@ def strategy():
 
 @pytest.fixture
 def mock_adapter_responses(strategy):
-    """Set up mock responses for adapter calls."""
     # Mock balance adapter
     strategy.balance_adapter.get_balance = AsyncMock(return_value=(True, 1000000))
     strategy.balance_adapter.move_from_main_wallet_to_strategy_wallet = AsyncMock(
@@ -119,7 +113,6 @@ def mock_adapter_responses(strategy):
 @pytest.mark.asyncio
 @pytest.mark.smoke
 async def test_smoke(strategy, mock_adapter_responses):
-    """REQUIRED: Basic smoke test - verifies strategy lifecycle."""
     examples = load_strategy_examples(Path(__file__))
     smoke_data = examples["smoke"]
 
@@ -152,7 +145,6 @@ async def test_smoke(strategy, mock_adapter_responses):
             assert isinstance(ok, bool)
             assert isinstance(msg, str)
 
-        # Update test
         with patch.object(strategy, "update", new_callable=AsyncMock) as mock_update:
             mock_update.return_value = (True, "success")
             ok, msg = await strategy.update(**smoke_data.get("update", {}))
@@ -169,7 +161,6 @@ async def test_smoke(strategy, mock_adapter_responses):
 
 @pytest.mark.asyncio
 async def test_canonical_usage(strategy, mock_adapter_responses):
-    """REQUIRED: Test canonical usage examples from examples.json (minimum)."""
     examples = load_strategy_examples(Path(__file__))
     canonical = get_canonical_examples(examples)
 
@@ -216,7 +207,6 @@ async def test_canonical_usage(strategy, mock_adapter_responses):
 
 @pytest.mark.asyncio
 async def test_status_returns_status_dict(strategy, mock_adapter_responses):
-    """Test that _status returns a proper StatusDict."""
     snap = MagicMock()
     snap.totals_usd = {}
     snap.ltv = 0.5
@@ -230,7 +220,7 @@ async def test_status_returns_status_dict(strategy, mock_adapter_responses):
         with patch.object(
             strategy, "_get_gas_balance", new_callable=AsyncMock
         ) as mock_gas:
-            mock_gas.return_value = 100000000000000000  # 0.1 ETH
+            mock_gas.return_value = 100000000000000000
             with patch.object(
                 strategy, "get_peg_diff", new_callable=AsyncMock
             ) as mock_peg:
@@ -251,7 +241,6 @@ async def test_status_returns_status_dict(strategy, mock_adapter_responses):
 
 @pytest.mark.asyncio
 async def test_policies_returns_list(strategy):
-    """Test that policies returns a non-empty list."""
     # Mock the policy functions to avoid ABI fetching
     with (
         patch(
@@ -292,7 +281,6 @@ async def test_policies_returns_list(strategy):
 
 @pytest.mark.asyncio
 async def test_quote_returns_apy_info(strategy, mock_adapter_responses):
-    """Test that quote returns APY information."""
     with patch("httpx.AsyncClient") as mock_client:
         mock_response = MagicMock()
         mock_response.json.return_value = {"data": {"smaApr": 3.5}}
@@ -310,8 +298,6 @@ async def test_quote_returns_apy_info(strategy, mock_adapter_responses):
 
 
 def test_max_safe_f_calculates_correctly(strategy):
-    """Test that _max_safe_F calculates the depeg-aware leverage limit correctly."""
-    # Set up strategy with MAX_DEPEG = 0.01 (1%)
     strategy.MAX_DEPEG = 0.01
 
     # With cf_w = 0.8, a = 0.99
@@ -322,7 +308,6 @@ def test_max_safe_f_calculates_correctly(strategy):
 
 
 def test_max_safe_f_with_zero_collateral_factor(strategy):
-    """Test _max_safe_F with zero collateral factor."""
     strategy.MAX_DEPEG = 0.01
     result = strategy._max_safe_F(0.0)
     # F_max = 1 / (1 + 0 * anything) = 1.0
@@ -333,15 +318,18 @@ def test_max_safe_f_with_zero_collateral_factor(strategy):
 async def test_swap_with_retries_succeeds_first_attempt(
     strategy, mock_adapter_responses
 ):
-    """Test that _swap_with_retries succeeds on first attempt."""
     strategy.max_swap_retries = 3
     strategy.swap_slippage_tolerance = 0.005
 
-    result = await strategy._swap_with_retries(
-        from_token_id="usd-coin-base",
-        to_token_id="l2-standard-bridged-weth-base-base",
-        amount=1000000,
-    )
+    with patch.object(
+        strategy, "_get_balance_raw", new_callable=AsyncMock
+    ) as mock_balance:
+        mock_balance.return_value = 10**18
+        result = await strategy._swap_with_retries(
+            from_token_id="usd-coin-base",
+            to_token_id="l2-standard-bridged-weth-base-base",
+            amount=1000000,
+        )
 
     assert result is not None
     assert "to_amount" in result
@@ -352,7 +340,6 @@ async def test_swap_with_retries_succeeds_first_attempt(
 async def test_swap_with_retries_succeeds_on_second_attempt(
     strategy, mock_adapter_responses
 ):
-    """Test that _swap_with_retries retries and succeeds."""
     strategy.max_swap_retries = 3
     strategy.swap_slippage_tolerance = 0.005
 
@@ -364,7 +351,13 @@ async def test_swap_with_retries_succeeds_on_second_attempt(
         ]
     )
 
-    with patch("asyncio.sleep", new_callable=AsyncMock):  # Speed up test
+    with (
+        patch.object(
+            strategy, "_get_balance_raw", new_callable=AsyncMock
+        ) as mock_balance,
+        patch("asyncio.sleep", new_callable=AsyncMock),
+    ):
+        mock_balance.return_value = 10**18
         result = await strategy._swap_with_retries(
             from_token_id="usd-coin-base",
             to_token_id="l2-standard-bridged-weth-base-base",
@@ -378,7 +371,6 @@ async def test_swap_with_retries_succeeds_on_second_attempt(
 
 @pytest.mark.asyncio
 async def test_swap_with_retries_fails_all_attempts(strategy, mock_adapter_responses):
-    """Test that _swap_with_retries returns None after all attempts fail."""
     strategy.max_swap_retries = 3
     strategy.swap_slippage_tolerance = 0.005
 
@@ -386,7 +378,13 @@ async def test_swap_with_retries_fails_all_attempts(strategy, mock_adapter_respo
         side_effect=Exception("Swap failed")
     )
 
-    with patch("asyncio.sleep", new_callable=AsyncMock):
+    with (
+        patch.object(
+            strategy, "_get_balance_raw", new_callable=AsyncMock
+        ) as mock_balance,
+        patch("asyncio.sleep", new_callable=AsyncMock),
+    ):
+        mock_balance.return_value = 10**18
         result = await strategy._swap_with_retries(
             from_token_id="usd-coin-base",
             to_token_id="l2-standard-bridged-weth-base-base",
@@ -401,7 +399,6 @@ async def test_swap_with_retries_fails_all_attempts(strategy, mock_adapter_respo
 async def test_swap_with_retries_aborts_on_unknown_outcome(
     strategy, mock_adapter_responses
 ):
-    """Swap retries must abort (no retry) when the transaction outcome is unknown."""
     strategy.max_swap_retries = 3
     strategy.brap_adapter.swap_from_token_ids = AsyncMock(
         return_value=(
@@ -410,19 +407,22 @@ async def test_swap_with_retries_aborts_on_unknown_outcome(
         )
     )
 
-    with pytest.raises(SwapOutcomeUnknownError):
-        await strategy._swap_with_retries(
-            from_token_id=USDC_TOKEN_ID,
-            to_token_id=WETH_TOKEN_ID,
-            amount=1000000,
-        )
+    with patch.object(
+        strategy, "_get_balance_raw", new_callable=AsyncMock
+    ) as mock_balance:
+        mock_balance.return_value = 10**18
+        with pytest.raises(SwapOutcomeUnknownError):
+            await strategy._swap_with_retries(
+                from_token_id=USDC_TOKEN_ID,
+                to_token_id=WETH_TOKEN_ID,
+                amount=1000000,
+            )
 
     assert strategy.brap_adapter.swap_from_token_ids.call_count == 1
 
 
 @pytest.mark.asyncio
 async def test_post_run_guard_no_action_when_delta_ok(strategy, mock_adapter_responses):
-    """Post-run guard should do nothing when delta is within tolerance."""
     snap = MagicMock()
     snap.wallet_wsteth = 0
     snap.usdc_supplied = 0
@@ -472,7 +472,6 @@ async def test_post_run_guard_no_action_when_delta_ok(strategy, mock_adapter_res
 async def test_post_run_guard_restores_delta_via_reconcile(
     strategy, mock_adapter_responses
 ):
-    """Post-run guard should attempt wallet reconcile in operate mode when net short."""
     snap1 = MagicMock()
     snap1.wallet_wsteth = 0
     snap1.usdc_supplied = 0
@@ -529,7 +528,6 @@ async def test_post_run_guard_restores_delta_via_reconcile(
 async def test_post_run_guard_delevers_to_delta_when_reconcile_insufficient(
     strategy, mock_adapter_responses
 ):
-    """Post-run guard should delever debt when still net short after reconcile."""
     snap1 = MagicMock()
     snap1.wallet_wsteth = 0
     snap1.usdc_supplied = 0
@@ -554,7 +552,7 @@ async def test_post_run_guard_delevers_to_delta_when_reconcile_insufficient(
     snap3.wallet_wsteth = 0
     snap3.usdc_supplied = 0
     snap3.wsteth_supplied = 0
-    snap3.weth_debt = int(0.005 * 10**18)  # small debt
+    snap3.weth_debt = int(0.005 * 10**18)
     snap3.debt_usd = 12.0
     snap3.hf = 1.3
     snap3.capacity_usd = 1000.0
@@ -602,7 +600,6 @@ async def test_post_run_guard_delevers_to_delta_when_reconcile_insufficient(
 async def test_atomic_deposit_iteration_swaps_from_eth_when_borrow_surfaces_as_eth(
     strategy, mock_adapter_responses
 ):
-    """Borrowed WETH can surface as native ETH; iteration should wrap ETH→WETH then swap WETH→wstETH."""
     # Ensure gas reserve exists so we don't drain to 0
     strategy.WRAP_GAS_RESERVE = 0.0014
 
@@ -615,12 +612,8 @@ async def test_atomic_deposit_iteration_swaps_from_eth_when_borrow_surfaces_as_e
         WSTETH_TOKEN_ID: 0,
     }
 
-    async def get_balance_side_effect(*, token_id: str, wallet_address: str, **_):
-        return (True, balances.get(token_id, 0))
-
-    strategy.balance_adapter.get_balance = AsyncMock(
-        side_effect=get_balance_side_effect
-    )
+    async def get_balance_raw_side_effect(*, token_id: str, wallet_address: str, **_):
+        return balances.get(token_id, 0)
 
     async def borrow_side_effect(*, mtoken: str, amount: int):
         # Borrow shows up as native ETH (simulates on-chain behavior)
@@ -650,7 +643,11 @@ async def test_atomic_deposit_iteration_swaps_from_eth_when_borrow_surfaces_as_e
 
     strategy._swap_with_retries = AsyncMock(side_effect=swap_side_effect)
 
-    lent = await strategy._atomic_deposit_iteration(borrow_amt_wei)
+    with patch.object(
+        strategy, "_get_balance_raw", new_callable=AsyncMock
+    ) as mock_balance:
+        mock_balance.side_effect = get_balance_raw_side_effect
+        lent = await strategy._atomic_deposit_iteration(borrow_amt_wei)
 
     assert lent == 123
     strategy.moonwell_adapter.borrow.assert_called_once_with(
@@ -664,46 +661,52 @@ async def test_atomic_deposit_iteration_swaps_from_eth_when_borrow_surfaces_as_e
 async def test_reconcile_wallet_into_position_uses_eth_inventory(
     strategy, mock_adapter_responses
 ):
-    """If debt exists but wstETH collateral is missing, use wallet ETH to complete the loop."""
-
-    async def get_pos_side_effect(*, mtoken: str):
-        if mtoken == M_USDC:
-            return (True, {"underlying_balance": 0, "borrow_balance": 0})
-        if mtoken == M_WSTETH:
-            return (True, {"underlying_balance": 0, "borrow_balance": 0})
-        if mtoken == M_WETH:
-            return (True, {"underlying_balance": 0, "borrow_balance": 5 * 10**18})
-        return (True, {})
-
-    strategy.moonwell_adapter.get_pos = AsyncMock(side_effect=get_pos_side_effect)
-
     balances: dict[str, int] = {
         ETH_TOKEN_ID: 10 * 10**18,
         WETH_TOKEN_ID: 0,
         WSTETH_TOKEN_ID: 0,
     }
 
-    async def get_balance_side_effect(*, token_id: str, wallet_address: str):
-        return (True, balances.get(token_id, 0))
+    snap = MagicMock()
+    snap.wallet_wsteth = 0
+    snap.wallet_weth = 0
+    snap.wsteth_dec = 18
+    snap.weth_dec = 18
+    snap.wsteth_price = 2000.0
+    snap.weth_price = 2000.0
+    snap.weth_debt = 5 * 10**18
+    snap.debt_usd = 10000.0
+    snap.eth_usable_wei = 10 * 10**18
+    snap.totals_usd = {f"Base_{M_WSTETH}": 0.0, f"Base_{WETH}": -10000.0}
 
-    strategy.balance_adapter.get_balance = AsyncMock(
-        side_effect=get_balance_side_effect
-    )
+    async def get_balance_raw_side_effect(*, token_id: str, wallet_address: str, **_):
+        return balances.get(token_id, 0)
 
     async def swap_side_effect(
         *, from_token_id: str, to_token_id: str, amount: int, **_
     ):
+        # _reconcile_wallet_into_position calls _swap_with_retries with ETH_TOKEN_ID
         assert from_token_id == ETH_TOKEN_ID
         assert to_token_id == WSTETH_TOKEN_ID
         balances[WSTETH_TOKEN_ID] += 7 * 10**18
-        return {"to_amount": 7 * 10**18}
+        return {"to_amount": 7 * 10**18, "block_number": 12345}
 
     strategy._swap_with_retries = AsyncMock(side_effect=swap_side_effect)
 
-    success, msg = await strategy._reconcile_wallet_into_position(
-        collateral_factors=(0.8, 0.8),
-        max_batch_usd=100000.0,
-    )
+    with (
+        patch.object(
+            strategy, "_accounting_snapshot", new_callable=AsyncMock
+        ) as mock_snap,
+        patch.object(
+            strategy, "_get_balance_raw", new_callable=AsyncMock
+        ) as mock_balance,
+    ):
+        mock_snap.return_value = (snap, (0.8, 0.8))
+        mock_balance.side_effect = get_balance_raw_side_effect
+        success, msg = await strategy._reconcile_wallet_into_position(
+            collateral_factors=(0.8, 0.8),
+            max_batch_usd=100000.0,
+        )
 
     assert success is True
     assert strategy._swap_with_retries.called
@@ -712,7 +715,6 @@ async def test_reconcile_wallet_into_position_uses_eth_inventory(
 
 @pytest.mark.asyncio
 async def test_sweep_token_balances_no_tokens(strategy, mock_adapter_responses):
-    """Test that _sweep_token_balances handles empty wallet."""
     # All balances are 0
     strategy.balance_adapter.get_balance = AsyncMock(return_value=(True, 0))
 
@@ -726,29 +728,30 @@ async def test_sweep_token_balances_no_tokens(strategy, mock_adapter_responses):
 
 @pytest.mark.asyncio
 async def test_sweep_token_balances_sweeps_tokens(strategy, mock_adapter_responses):
-    """Test that _sweep_token_balances converts dust tokens."""
     strategy.min_withdraw_usd = 1.0
 
     # Mock balance returns (has some WETH dust)
-    def balance_side_effect(token_id, wallet_address):
+    async def get_balance_raw_side_effect(*, token_id: str, wallet_address: str, **_):
         if "weth" in token_id.lower():
-            return (True, 100 * 10**18)  # 100 WETH
-        return (True, 0)
-
-    strategy.balance_adapter.get_balance = AsyncMock(side_effect=balance_side_effect)
+            return 100 * 10**18
+        return 0
 
     # Mock price (high enough to trigger sweep)
     strategy.token_adapter.get_token_price = AsyncMock(
         return_value=(True, {"current_price": 2000.0})
     )
 
-    success, msg = await strategy._sweep_token_balances(
-        target_token_id="usd-coin-base",
-        exclude=set(),
-    )
+    with patch.object(
+        strategy, "_get_balance_raw", new_callable=AsyncMock
+    ) as mock_balance:
+        mock_balance.side_effect = get_balance_raw_side_effect
+        success, msg = await strategy._sweep_token_balances(
+            target_token_id="usd-coin-base",
+            exclude=set(),
+        )
 
     assert success is True
-    # Should have called swap
+    # Should have called swap via _swap_with_retries
     strategy.brap_adapter.swap_from_token_ids.assert_called()
 
 
@@ -757,7 +760,6 @@ async def test_sweep_token_balances_sweeps_tokens(strategy, mock_adapter_respons
 
 @pytest.mark.asyncio
 async def test_deposit_rejects_zero_amount(strategy):
-    """Test that deposit rejects zero or negative amounts."""
     result = await strategy.deposit(main_token_amount=0.0)
     assert result[0] is False
     assert "positive" in result[1].lower()
@@ -768,26 +770,22 @@ async def test_deposit_rejects_zero_amount(strategy):
 
 
 def test_slippage_capped_at_max(strategy):
-    """Test that slippage is capped at MAX_SLIPPAGE_TOLERANCE."""
     strategy.MAX_SLIPPAGE_TOLERANCE = 0.03
     strategy.swap_slippage_tolerance = 0.02
 
     # With 3 retries at 2% base: 2%, 4%, 6% -> should be capped at 3%
-    # The actual slippage calculation happens in the method, we just verify the constant exists
     assert hasattr(strategy, "MAX_SLIPPAGE_TOLERANCE")
     assert strategy.MAX_SLIPPAGE_TOLERANCE == 0.03
 
 
 def test_price_staleness_threshold_exists(strategy):
-    """Test that price staleness threshold is configured."""
     assert hasattr(strategy, "PRICE_STALENESS_THRESHOLD")
     assert strategy.PRICE_STALENESS_THRESHOLD > 0
 
 
 def test_min_leverage_gain_constant_exists(strategy):
-    """Test that minimum leverage gain constant is configured."""
     assert hasattr(strategy, "_MIN_LEVERAGE_GAIN_BPS")
-    assert strategy._MIN_LEVERAGE_GAIN_BPS == 50e-4  # 50 bps
+    assert strategy._MIN_LEVERAGE_GAIN_BPS == 50e-4
 
 
 @pytest.mark.asyncio
@@ -848,7 +846,6 @@ async def test_withdraw_runs_post_run_guard_only_on_failure(
 
 @pytest.mark.asyncio
 async def test_leverage_calc_handles_high_cf_w(strategy, mock_adapter_responses):
-    """Test that leverage calculation handles high collateral factors safely."""
     strategy.MIN_HEALTH_FACTOR = 1.2
 
     # This should return early without crashing when cf_w >= MIN_HEALTH_FACTOR
@@ -861,20 +858,19 @@ async def test_leverage_calc_handles_high_cf_w(strategy, mock_adapter_responses)
         initial_leverage=1.5,
         usdc_lend_value=1000.0,
         wsteth_lend_value=500.0,
-        collateral_factors=(0.8, 1.3),  # cf_u=0.8, cf_w=1.3 (higher than MIN_HF)
+        collateral_factors=(0.8, 1.3),
     )
 
     # Should return failure tuple instead of crashing
     assert result[0] is False
-    assert result[2] == -1  # Error code
+    assert result[2] == -1
 
 
 @pytest.mark.asyncio
 async def test_price_staleness_triggers_refresh(strategy, mock_adapter_responses):
-    """Test that stale prices trigger a refresh."""
-    strategy.PRICE_STALENESS_THRESHOLD = 1  # 1 second for test
+    strategy.PRICE_STALENESS_THRESHOLD = 1
     strategy._token_price_cache = {"test-token": 100.0}
-    strategy._token_price_timestamps = {"test-token": 0}  # Very old timestamp
+    strategy._token_price_timestamps = {"test-token": 0}
 
     # Should refresh because timestamp is stale
     await strategy._get_token_price("test-token")
@@ -885,8 +881,6 @@ async def test_price_staleness_triggers_refresh(strategy, mock_adapter_responses
 
 @pytest.mark.asyncio
 async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
-    """Partial liquidation should redeem wstETH first when collateral exceeds debt."""
-
     # Token metadata
     async def mock_get_token(token_id: str):
         if token_id == USDC_TOKEN_ID:
@@ -906,10 +900,8 @@ async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
     # Wallet balances (raw)
     balances: dict[str, int] = {USDC_TOKEN_ID: 0, WSTETH_TOKEN_ID: 0}
 
-    async def mock_get_balance(*, token_id: str, wallet_address: str):
-        return (True, balances.get(token_id, 0))
-
-    strategy.balance_adapter.get_balance = AsyncMock(side_effect=mock_get_balance)
+    async def mock_get_balance_raw(*, token_id: str, wallet_address: str, **_):
+        return balances.get(token_id, 0)
 
     # Position snapshot: wstETH collateral > WETH debt
     totals_usd = {
@@ -923,8 +915,6 @@ async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
     snap.wsteth_price = 2000.0
     snap.wsteth_dec = 18
 
-    strategy._accounting_snapshot = AsyncMock(return_value=(snap, (0.8, 0.8)))
-
     # Collateral factors
     strategy.moonwell_adapter.get_collateral_factor = AsyncMock(
         return_value=(True, 0.8)
@@ -948,7 +938,7 @@ async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
     async def mock_unlend(*, mtoken: str, amount: int):
         if mtoken == M_WSTETH:
             balances[WSTETH_TOKEN_ID] += int(amount)
-        return (True, "success")
+        return (True, {"block_number": 12345})
 
     strategy.moonwell_adapter.unlend = AsyncMock(side_effect=mock_unlend)
 
@@ -967,14 +957,25 @@ async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
         usd_out = (amt / 10**18) * 2000.0
         usdc_out = int(usd_out * 10**6)
         balances[USDC_TOKEN_ID] += usdc_out
-        return (True, {"to_amount": usdc_out})
+        return (True, {"to_amount": usdc_out, "block_number": 12346})
 
     strategy.brap_adapter.swap_from_token_ids = AsyncMock(side_effect=mock_swap)
 
     # Also need to mock lend since partial_liquidate may try to re-lend leftover wstETH
     strategy.moonwell_adapter.lend = AsyncMock(return_value=(True, "success"))
 
-    ok, msg = await strategy.partial_liquidate(usd_value=100.0)
+    with (
+        patch.object(
+            strategy, "_accounting_snapshot", new_callable=AsyncMock
+        ) as mock_snap,
+        patch.object(
+            strategy, "_get_balance_raw", new_callable=AsyncMock
+        ) as mock_balance,
+    ):
+        mock_snap.return_value = (snap, (0.8, 0.8))
+        mock_balance.side_effect = mock_get_balance_raw
+        ok, msg = await strategy.partial_liquidate(usd_value=100.0)
+
     assert ok
     assert "available" in msg.lower()
 
@@ -986,7 +987,6 @@ async def test_partial_liquidate_prefers_wsteth_when_excess(strategy):
 
 @pytest.mark.asyncio
 async def test_partial_liquidate_uses_usdc_collateral_when_no_wsteth_excess(strategy):
-    """If wstETH collateral doesn't exceed debt, partial liquidation should redeem USDC collateral."""
     # Token metadata
     strategy.token_adapter.get_token = AsyncMock(
         side_effect=lambda token_id: (
@@ -1000,13 +1000,11 @@ async def test_partial_liquidate_uses_usdc_collateral_when_no_wsteth_excess(stra
 
     balances: dict[str, int] = {USDC_TOKEN_ID: 0}
 
-    async def mock_get_balance(*, token_id: str, wallet_address: str):
-        return (True, balances.get(token_id, 0))
-
-    strategy.balance_adapter.get_balance = AsyncMock(side_effect=mock_get_balance)
+    async def mock_get_balance_raw(*, token_id: str, wallet_address: str, **_):
+        return balances.get(token_id, 0)
 
     totals_usd = {
-        f"Base_{M_WSTETH}": 100.0,  # <= debt
+        f"Base_{M_WSTETH}": 100.0,
         f"Base_{M_USDC}": 500.0,
         f"Base_{WETH}": -200.0,
     }
@@ -1016,8 +1014,6 @@ async def test_partial_liquidate_uses_usdc_collateral_when_no_wsteth_excess(stra
     snap.usdc_price = 1.0
     snap.usdc_dec = 6
 
-    strategy._accounting_snapshot = AsyncMock(return_value=(snap, (0.8, 0.8)))
-
     strategy.moonwell_adapter.get_collateral_factor = AsyncMock(
         return_value=(True, 0.8)
     )
@@ -1039,11 +1035,22 @@ async def test_partial_liquidate_uses_usdc_collateral_when_no_wsteth_excess(stra
     async def mock_unlend(*, mtoken: str, amount: int):
         if mtoken == M_USDC:
             balances[USDC_TOKEN_ID] += int(amount)
-        return (True, "success")
+        return (True, {"block_number": 12345})
 
     strategy.moonwell_adapter.unlend = AsyncMock(side_effect=mock_unlend)
 
-    ok, msg = await strategy.partial_liquidate(usd_value=50.0)
+    with (
+        patch.object(
+            strategy, "_accounting_snapshot", new_callable=AsyncMock
+        ) as mock_snap,
+        patch.object(
+            strategy, "_get_balance_raw", new_callable=AsyncMock
+        ) as mock_balance,
+    ):
+        mock_snap.return_value = (snap, (0.8, 0.8))
+        mock_balance.side_effect = mock_get_balance_raw
+        ok, msg = await strategy.partial_liquidate(usd_value=50.0)
+
     assert ok
     assert "available" in msg.lower()