new-site/scripts/workers/crypto_offramp/mock.py

"""In-memory mock offramp for tests.

Same method surface as CoinbasePrimeOfframp but never hits the network.
Records calls for assertion. State advances deterministically based on
``auto_complete`` (default True) or manually via ``.advance()``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from decimal import Decimal
from typing import Optional
from uuid import uuid4

from . import (
    ExecutionResult,
    KycHoldError,
    OfframpError,
    ProviderStatus,
    Quote,
    Rail,
    SUPPORTED_COINS,
)


@dataclass
class _MockTransfer:
    provider_ref: str
    amount_usd_cents: int
    rail: Rail
    state: str = "submitted"
    memo: Optional[str] = None


class MockOfframp:
    name = "mock"
    accepts_coins = SUPPORTED_COINS
    supports_rails: set[Rail] = {"ach", "wire", "rtp"}

    def __init__(self, *, auto_complete: bool = True, kyc_hold_for_coins: Optional[set[str]] = None):
        self.auto_complete = auto_complete
        self.kyc_hold_for_coins = kyc_hold_for_coins or set()
        self.calls: list[tuple[str, dict]] = []
        self._transfers: dict[str, _MockTransfer] = {}
        # Fixed synthetic prices
        self._prices = {
            "BTC":   Decimal("70000.00"),
            "ETH":   Decimal("3500.00"),
            "MATIC": Decimal("0.80"),
            "LTC":   Decimal("85.00"),
            "TRX":   Decimal("0.11"),
            "BNB":   Decimal("520.00"),
            "DOGE":  Decimal("0.15"),
            "USDC":  Decimal("1.00"),
        }

    async def quote(self, amount_coin: Decimal, from_coin: str) -> Quote:
        self.calls.append(("quote", {"amount_coin": amount_coin, "from_coin": from_coin}))
        fx = self._prices.get(from_coin.upper(), Decimal("1.00"))
        gross = amount_coin * fx
        fee = gross * Decimal("0.0008")
        return Quote(
            from_coin=from_coin.upper(),
            amount_coin=amount_coin,
            estimated_usd_cents=int((gross - fee) * 100),
            fx_rate_usd=fx,
            fee_usd_cents=int(fee * 100),
        )

    async def deposit_address(self, coin: str) -> str:
        self.calls.append(("deposit_address", {"coin": coin}))
        return f"mock-deposit-{coin.lower()}-{uuid4().hex[:8]}"

    async def execute(
        self, *,
        amount_coin: Decimal, from_coin: str, rail: Rail,
        idempotency_key: str,
        relay_routing: Optional[str] = None,
        relay_account: Optional[str] = None,
        memo: Optional[str] = None,
        target_usd_cents: Optional[int] = None,
    ) -> ExecutionResult:
        self.calls.append(("execute", {
            "amount_coin": amount_coin, "from_coin": from_coin, "rail": rail,
            "idempotency_key": idempotency_key, "memo": memo,
            "target_usd_cents": target_usd_cents,
        }))
        if from_coin.upper() in self.kyc_hold_for_coins:
            raise KycHoldError(f"mock KYC hold for {from_coin}")
        fx = self._prices.get(from_coin.upper(), Decimal("1.00"))
        usd_cents = target_usd_cents or int(amount_coin * fx * 100)
        ref = f"mock-tx-{uuid4().hex[:12]}"
        state = "completed" if self.auto_complete else "submitted"
        self._transfers[ref] = _MockTransfer(
            provider_ref=ref, amount_usd_cents=usd_cents, rail=rail,
            state=state, memo=memo,
        )
        return ExecutionResult(
            provider_ref=ref,
            accepted_usd_cents=usd_cents,
            fill_price_usd=fx,
            state=state,
            rail=rail,
            memo=memo,
            raw={"mock": True},
        )

    async def status(self, provider_ref: str) -> ProviderStatus:
        self.calls.append(("status", {"provider_ref": provider_ref}))
        t = self._transfers.get(provider_ref)
        if not t:
            raise OfframpError(f"mock: no transfer {provider_ref}")
        return ProviderStatus(
            provider_ref=provider_ref,
            state=t.state,  # type: ignore[arg-type]
            settled_usd_cents=t.amount_usd_cents if t.state == "completed" else 0,
        )

    # ── Test helpers ────────────────────────────────────────────────────

    def advance(self, provider_ref: str, to_state: str):
        """Manually move a transfer forward in state."""
        if provider_ref in self._transfers:
            self._transfers[provider_ref].state = to_state


# ── Mock SHKeeper client ────────────────────────────────────────────────


@dataclass
class _MockPayout:
    withdraw_id: str
    tx_hash: Optional[str]
    fee_coin: Decimal = Decimal("0")
    raw: Optional[dict] = None


class MockSHKeeperClient:
    """In-memory SHKeeperClient substitute for tests — never hits the network."""
    def __init__(self, *, default_balances: Optional[dict[str, Decimal]] = None):
        self._balances: dict[str, Decimal] = dict(default_balances or {})
        self.payouts: list[tuple[str, str, Decimal]] = []

    async def payout(self, *, coin: str, destination: str, amount: Decimal):
        self.payouts.append((coin, destination, amount))
        self._balances[coin] = self._balances.get(coin, Decimal("0")) - amount
        return _MockPayout(
            withdraw_id=f"mock-shk-{uuid4().hex[:10]}",
            tx_hash=f"mock-txhash-{uuid4().hex[:16]}",
        )

    async def balance(self, coin: str):
        from .shkeeper_client import BalanceResult
        bal = self._balances.get(coin.upper(), Decimal("0"))
        return BalanceResult(balance_coin=bal, balance_usd=bal * Decimal("70000"))

    async def withdrawal_status(self, coin: str, withdraw_id: str):
        return {"status": "confirmed", "txid": f"mock-tx-{withdraw_id}"}