"""OKX 资金账户余额(asset/balances)."""
from __future__ import annotations

import unittest
from unittest.mock import MagicMock

from lib.exchange.okx_options_lib import (
    fetch_funding_balances_via_asset_api,
    fetch_options_balances,
)


class TestOkxFundingBalances(unittest.TestCase):
    def test_asset_balances_parsed(self):
        ex = MagicMock()
        ex.private_get_asset_balances.return_value = {
            "data": [
                {"ccy": "USDT", "availBal": "25.5", "bal": "30"},
                {"ccy": "USDC", "availBal": "10", "bal": "10"},
            ]
        }
        total, avail = fetch_funding_balances_via_asset_api(ex)
        self.assertEqual(avail["USDT"], 25.5)
        self.assertEqual(total["USDT"], 30.0)
        self.assertEqual(avail["USDC"], 10.0)

    def test_fetch_options_balances_merges_asset_api(self):
        ex = MagicMock()
        ex.fetch_balance.return_value = {"free": {}, "total": {}}
        ex.private_get_asset_balances.return_value = {
            "data": [{"ccy": "USDT", "availBal": "18.2", "bal": "18.2"}]
        }
        bal = fetch_options_balances(ex, force=True)
        self.assertEqual(bal["funding_usdt_avail"], 18.2)
        self.assertEqual(bal["funding_usdt"], 18.2)


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