trex-core/tests/domain/test_capital_allocation.py

"""Shared account-capital allocation policy regression tests."""
from __future__ import annotations

import random
from decimal import Decimal

import pytest

from trex_core.domain.capital_allocation import (
    CapitalCandidate,
    allocate_liquid_slots,
    plan_liquid_capital,
    protected_share_cost,
)


def candidate(
    portfolio_id: str,
    symbol: str,
    cost: int,
) -> CapitalCandidate:
    return CapitalCandidate(portfolio_id, symbol, cost)


def test_mega_two_million_excludes_hynix_and_uses_samsung_lane() -> None:
    plan = plan_liquid_capital(
        2_000_000,
        slot_count=2,
        candidates=(
            candidate("samsung", "005930", 220_000),
            candidate("hynix", "000660", 1_400_000),
        ),
    )

    assert plan.caps == {"samsung": 2_000_000, "hynix": 0}
    assert plan.allocations[1].reason == "protected_share_above_flex_limit"


def test_mega_three_million_flexes_one_share_into_hynix_lane() -> None:
    assert allocate_liquid_slots(
        3_000_000,
        slot_count=2,
        candidates=(
            candidate("samsung", "005930", 220_000),
            candidate("hynix", "000660", 1_510_000),
        ),
    ) == {"samsung": 1_490_000, "hynix": 1_510_000}


def test_top50_three_slots_exclude_unaffordable_high_price_candidate() -> None:
    plan = plan_liquid_capital(
        2_000_000,
        slot_count=3,
        candidates=(
            candidate("rank1", "000001", 100_000),
            candidate("rank2", "000002", 200_000),
            candidate("rank3", "000003", 500_000),
            candidate("rank4", "000004", 700_000),
            candidate("rank5", "000005", 1_050_000),
        ),
    )

    assert plan.caps == {
        "rank1": 650_000,
        "rank2": 650_000,
        "rank3": 650_000,
        "rank4": 700_000,
        "rank5": 0,
    }
    assert plan.allocations[-1].reason == "protected_share_above_flex_limit"


def test_zero_cash_fails_closed_with_explicit_reason() -> None:
    plan = plan_liquid_capital(
        0,
        slot_count=3,
        candidates=(candidate("a", "000001", 100_000),),
    )

    assert plan.caps == {"a": 0}
    assert plan.allocations[0].reason == "no_available_cash"


def test_protected_share_cost_includes_execution_margin() -> None:
    assert protected_share_cost(
        Decimal("1433000"),
        Decimal("0.5"),
    ) == 1_440_165


@pytest.mark.parametrize(
    ("price", "slippage"),
    [
        (Decimal("0"), Decimal("0.5")),
        (Decimal("-1"), Decimal("0.5")),
        (Decimal("NaN"), Decimal("0.5")),
        (Decimal("1"), Decimal("-0.1")),
        (Decimal("1"), Decimal("Infinity")),
    ],
)
def test_protected_share_cost_rejects_invalid_inputs(
    price: Decimal,
    slippage: Decimal,
) -> None:
    with pytest.raises(ValueError, match="finite"):
        protected_share_cost(price, slippage)


@pytest.mark.parametrize(
    "kwargs,error",
    [
        (
            {
                "available_cash": -1,
                "slot_count": 2,
                "candidates": (candidate("a", "000001", 1),),
            },
            "cash",
        ),
        (
            {
                "available_cash": 1,
                "slot_count": 0,
                "candidates": (candidate("a", "000001", 1),),
            },
            "slot",
        ),
        (
            {
                "available_cash": 1,
                "slot_count": 1,
                "candidates": (),
            },
            "candidates",
        ),
    ],
)
def test_invalid_top_level_inputs_are_rejected(kwargs, error) -> None:
    with pytest.raises(ValueError, match=error):
        plan_liquid_capital(**kwargs)


def test_duplicate_identity_is_rejected() -> None:
    with pytest.raises(ValueError, match="must be unique"):
        plan_liquid_capital(
            1_000_000,
            slot_count=2,
            candidates=(
                candidate("same", "000001", 100_000),
                candidate("same", "000002", 100_000),
            ),
        )


def test_invalid_candidate_cost_and_flex_are_rejected() -> None:
    with pytest.raises(ValueError, match="protected share costs"):
        plan_liquid_capital(
            1_000_000,
            slot_count=1,
            candidates=(candidate("a", "000001", 0),),
        )
    with pytest.raises(ValueError, match="flex share"):
        plan_liquid_capital(
            1_000_000,
            slot_count=1,
            candidates=(candidate("a", "000001", 1),),
            flex_share_count=-1,
        )


def test_randomized_worst_active_combination_never_exceeds_cash() -> None:
    rng = random.Random(20260729)
    for case in range(10_000):
        slot_count = rng.randint(1, 5)
        candidate_count = rng.randint(1, 8)
        cash = rng.randint(0, 10_000_000)
        candidates = tuple(
            candidate(
                f"portfolio-{index}",
                f"{index:06d}",
                rng.randint(1, 3_000_000),
            )
            for index in range(1, candidate_count + 1)
        )
        plan = plan_liquid_capital(
            cash,
            slot_count=slot_count,
            candidates=candidates,
            flex_share_count=rng.randint(0, 2),
        )
        positive = sorted(
            (cap for cap in plan.caps.values() if cap > 0),
            reverse=True,
        )
        assert sum(positive[:slot_count]) <= cash, case
        for item, source in zip(plan.allocations, candidates):
            assert item.cap == 0 or item.cap >= source.protected_share_cost