crypto_monitor/tests/test_hub_calculator_lib.py

"""hub_calculator_lib 测算逻辑。"""

import unittest
from unittest.mock import patch

from hub_calculator_lib import (
    calc_initial_roll_qty,
    calc_roll_calculator,
    calc_trend_calculator,
    solve_add_amount_for_total_risk,
)

MOCK_MARKET = {
    "exchange_id": "0",
    "exchange_key": "binance",
    "exchange_name": "币安 · crypto_monitor_binance",
    "exchange_label": "币安 · crypto_monitor_binance",
    "base": "ETH",
    "exchange_symbol": "ETH/USDT:USDT",
    "display_symbol": "ETH/USDT",
    "contract_size": 1.0,
    "price_tick": 0.01,
    "price_decimals": 2,
    "amount_decimals": 3,
    "min_amount": 0.001,
}


def _mock_resolve(_exchange="binance", _base="ETH"):
    return MOCK_MARKET, lambda amount: round(float(amount), 3), None


class HubCalculatorLibTests(unittest.TestCase):
    @patch("hub_calculator_lib._resolve_market", return_value=_mock_resolve())
    def test_trend_calculator_long_basic(self, _mock):
        data, err = calc_trend_calculator(
            direction="long",
            capital_usdt=1000,
            risk_percent=5,
            leverage=5,
            entry_price=100,
            stop_loss=95,
            add_upper=110,
            take_profit=120,
            dca_legs=3,
            exchange_id="0",
            base="ETH",
        )
        self.assertIsNone(err)
        self.assertIsNotNone(data)
        assert data is not None
        self.assertEqual(data["risk_budget_u"], 50.0)
        self.assertGreaterEqual(len(data["rows"]), 2)
        self.assertEqual(data["rows"][0]["label"], "首仓")
        self.assertEqual(data["market"]["display_symbol"], "ETH/USDT")

    @patch("hub_calculator_lib._resolve_market", return_value=_mock_resolve())
    def test_trend_calculator_short_rejects_bad_bounds(self, _mock):
        data, err = calc_trend_calculator(
            direction="short",
            capital_usdt=1000,
            risk_percent=5,
            leverage=5,
            entry_price=100,
            stop_loss=90,
            add_upper=110,
            take_profit=80,
            dca_legs=3,
        )
        self.assertIsNone(data)
        self.assertIsNotNone(err)

    @patch("hub_calculator_lib._resolve_market", return_value=_mock_resolve())
    def test_roll_calculator_first_leg_auto(self, _mock):
        data, err = calc_roll_calculator(
            direction="long",
            capital_usdt=1000,
            risk_percent=5,
            entry_price=100,
            stop_loss=95,
            take_profit=120,
            add_legs=[],
            legs_done=0,
        )
        self.assertIsNone(err)
        self.assertIsNotNone(data)
        assert data is not None
        self.assertEqual(data["first_contracts"], 10.0)
        self.assertEqual(len(data["rows"]), 1)
        self.assertEqual(data["rows"][0]["loss_at_sl_u"], 50.0)
        self.assertEqual(data["rows"][0]["profit_at_tp_u"], 200.0)

    @patch("hub_calculator_lib._resolve_market", return_value=_mock_resolve())
    def test_roll_calculator_chain_two_legs(self, _mock):
        data, err = calc_roll_calculator(
            direction="long",
            capital_usdt=1000,
            risk_percent=5,
            entry_price=100,
            stop_loss=95,
            take_profit=120,
            add_legs=[
                {"add_price": 105, "new_stop_loss": 98},
                {"add_price": 108, "new_stop_loss": 101},
            ],
            legs_done=0,
        )
        self.assertIsNone(err)
        self.assertIsNotNone(data)
        assert data is not None
        self.assertEqual(len(data["rows"]), 3)
        self.assertEqual(data["rows"][1]["label"], "滚仓1")
        self.assertGreater(float(data["final_contracts"]), float(data["first_contracts"]))

    @patch("hub_calculator_lib._resolve_market", return_value=_mock_resolve())
    def test_roll_calculator_rejects_too_many_legs(self, _mock):
        data, err = calc_roll_calculator(
            direction="long",
            capital_usdt=1000,
            risk_percent=5,
            entry_price=100,
            stop_loss=95,
            take_profit=120,
            add_legs=[
                {"add_price": 105, "new_stop_loss": 98},
                {"add_price": 108, "new_stop_loss": 101},
                {"add_price": 110, "new_stop_loss": 103},
                {"add_price": 112, "new_stop_loss": 105},
            ],
            legs_done=0,
        )
        self.assertIsNone(data)
        self.assertIsNotNone(err)

    def test_initial_roll_qty(self):
        qty, err = calc_initial_roll_qty("long", 100, 95, 50, 1.0)
        self.assertIsNone(err)
        self.assertEqual(qty, 10.0)

    def test_initial_roll_qty_with_contract_size(self):
        qty, err = calc_initial_roll_qty("long", 100, 95, 50, 0.1)
        self.assertIsNone(err)
        self.assertEqual(qty, 100.0)

    def test_solve_add_with_contract_size(self):
        q2, err = solve_add_amount_for_total_risk(
            "long",
            qty_existing=10.0,
            entry_existing=100.0,
            add_price=105.0,
            new_stop=98.0,
            risk_budget_usdt=50.0,
            contract_size=1.0,
        )
        self.assertIsNone(err)
        self.assertIsNotNone(q2)
        assert q2 is not None
        self.assertGreater(q2, 0)


if __name__ == "__main__":
    unittest.main()