trex-core/tests/application/test_account_snapshot.py

"""Bounded account-read recovery (dcy-041 / GitHub #45)."""
from __future__ import annotations

from datetime import UTC, datetime, timedelta

import pytest

from trex_core.application.account_snapshot import (
    AccountSnapshot,
    AccountSnapshotService,
    AccountSnapshotUnavailable,
    InMemoryAccountSnapshotStore,
)
from trex_core.application.broker_port import BalanceItem, BrokerBalanceError


class FakeAccountBroker:
    def __init__(self) -> None:
        self.account_hash = "a" * 64
        self.balance = [BalanceItem("005930", 2)]
        self.cash = 500_000
        self.balance_error: str | None = None
        self.cash_error: str | None = None

    def account_scope_hash(self) -> str:
        return self.account_hash

    def get_balance(self) -> list[BalanceItem]:
        if self.balance_error is not None:
            raise BrokerBalanceError(self.balance_error)
        return list(self.balance)

    def get_available_cash(self, *, price_check_code: str = "005930") -> int:
        if self.cash_error is not None:
            raise BrokerBalanceError(self.cash_error)
        return self.cash


class BrokenLoadStore(InMemoryAccountSnapshotStore):
    def load(self, account_hash: str) -> AccountSnapshot | None:
        raise RuntimeError("db read failed")


def _wire():
    now = [datetime(2026, 7, 29, tzinfo=UTC)]
    broker = FakeAccountBroker()
    store = InMemoryAccountSnapshotStore()
    service = AccountSnapshotService(
        broker=broker,
        store=store,
        buy_max_age_sec=600,
        sell_max_age_sec=3600,
        clock=lambda: now[0],
        rng=lambda: 0.5,
    )
    return broker, store, service, now


def test_refresh_persists_independent_holdings_and_cash_capture_times():
    broker, store, service, now = _wire()
    assert service.refresh() is True
    first = store.load(broker.account_hash)
    assert first is not None

    now[0] += timedelta(seconds=120)
    broker.balance = [BalanceItem("005930", 3)]
    broker.cash_error = "transport_failed"
    assert service.refresh() is False

    second = store.load(broker.account_hash)
    assert second is not None
    assert second.holdings == (BalanceItem("005930", 3),)
    assert second.holdings_captured_at == now[0]
    assert second.available_cash == first.available_cash
    assert second.cash_captured_at == first.cash_captured_at


def test_sell_falls_back_only_for_transport_failure_and_within_sell_ttl():
    broker, _, service, now = _wire()
    assert service.refresh()
    broker.balance_error = "transport_failed"
    now[0] += timedelta(seconds=601)

    assert service.holdings_for_sell() == [BalanceItem("005930", 2)]
    assert service.buy_guard(qty=1, limit_price=100_000) == (
        "account_snapshot_unavailable"
    )

    now[0] += timedelta(seconds=3000)
    with pytest.raises(AccountSnapshotUnavailable) as raised:
        service.holdings_for_sell()
    assert raised.value.code == "expired"


@pytest.mark.parametrize(
    "error_code",
    ["authentication_failed", "authorization_failed", "malformed_response",
     "broker_rejected", "ownership_mismatch"],
)
def test_sell_never_falls_back_for_non_transport_failures(error_code):
    broker, _, service, _ = _wire()
    assert service.refresh()
    broker.balance_error = error_code

    with pytest.raises(AccountSnapshotUnavailable) as raised:
        service.holdings_for_sell()
    assert raised.value.code == "live_read_failed"


def test_buy_guard_accounts_for_existing_pending_cash_and_never_renews_age():
    broker, store, service, now = _wire()
    assert service.refresh()
    captured = store.load(broker.account_hash).cash_captured_at

    assert service.buy_guard(
        qty=2, limit_price=100_000, reserved_cash=250_000
    ) is None
    assert service.buy_guard(
        qty=3, limit_price=100_000, reserved_cash=250_000
    ) == "insufficient_snapshot_cash"
    assert store.load(broker.account_hash).cash_captured_at == captured

    now[0] += timedelta(seconds=601)
    assert service.buy_guard(qty=1, limit_price=1) == (
        "account_snapshot_unavailable"
    )


def test_current_account_hash_is_required_on_every_consume():
    broker, _, service, _ = _wire()
    assert service.refresh()
    broker.account_hash = "b" * 64

    assert service.buy_guard(qty=1, limit_price=1) == (
        "account_snapshot_unavailable"
    )
    broker.balance_error = "transport_failed"
    with pytest.raises(AccountSnapshotUnavailable) as raised:
        service.holdings_for_sell()
    assert raised.value.code == "account_mismatch"


def test_restart_transport_attempt_can_use_existing_fresh_cash_snapshot():
    broker, store, service, now = _wire()
    store.save(
        AccountSnapshot(
            account_hash=broker.account_hash,
            holdings=(BalanceItem("005930", 2),),
            holdings_captured_at=now[0],
            available_cash=500_000,
            cash_captured_at=now[0],
        )
    )
    broker.balance_error = "transport_failed"
    broker.cash_error = "transport_failed"

    assert service.refresh() is False
    assert service.buy_guard(qty=2, limit_price=100_000) is None


def test_invalid_or_missing_buy_notional_fails_closed():
    _, _, service, _ = _wire()
    assert service.refresh()
    assert service.buy_guard(qty=None, limit_price=1) == (
        "account_snapshot_price_required"
    )
    assert service.buy_guard(qty=1, limit_price=None) == (
        "account_snapshot_price_required"
    )


def test_slots_rpc_uses_snapshot_only_with_sell_ttl_freshness():
    broker, _, service, now = _wire()
    assert service.refresh()
    broker.balance = [BalanceItem("005930", 999)]

    holdings, captured_at = service.holdings_for_slots_rpc()
    assert holdings == [BalanceItem("005930", 2)]
    assert captured_at == now[0]

    now[0] += timedelta(seconds=3601)
    with pytest.raises(AccountSnapshotUnavailable) as raised:
        service.holdings_for_slots_rpc()
    assert raised.value.code == "expired"


def test_slots_rpc_normalizes_snapshot_store_read_failure():
    broker = FakeAccountBroker()
    service = AccountSnapshotService(
        broker=broker,
        store=BrokenLoadStore(),
        buy_max_age_sec=600,
        sell_max_age_sec=3600,
        clock=lambda: datetime(2026, 7, 29, tzinfo=UTC),
        rng=lambda: 0.5,
    )

    with pytest.raises(AccountSnapshotUnavailable) as raised:
        service.holdings_for_slots_rpc()
    assert raised.value.code == "read_failed"