stock-ai-agent/backend/tests/test_bitget_balance.py
2026-03-30 08:30:36 +08:00

357 lines
13 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
BitgetTradingAPI.get_balance() 单元测试
覆盖场景:
1. fetch_balance 成功 → 直接返回(不调原始 API
2. fetch_balance 无 USDT → 回退到原始 API (UTA)
3. fetch_balance 无 USDT → 回退到原始 API (合约)
4. fetch_balance 抛异常 → 回退到原始 API (UTA)
5. fetch_balance 和原始 API 都失败 → 返回 {}
6. 余额为 0 / None 的边界情况
7. get_account_state 端到端解析
"""
import importlib.util
import os
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
def load_bitget_sdk_class():
sdk_path = (
Path(__file__).resolve().parents[1]
/ 'app' / 'services' / 'bitget_trading_api_sdk.py'
)
if 'ccxt' not in sys.modules:
ccxt_module = types.ModuleType('ccxt')
ccxt_module.BaseError = Exception
ccxt_module.bitget = MagicMock()
sys.modules['ccxt'] = ccxt_module
# mock app.config
if 'app.config' not in sys.modules:
mock_cfg = types.ModuleType('app.config')
mock_cfg.get_settings = MagicMock()
sys.modules['app.config'] = mock_cfg
# mock app.utils.logger
if 'app.utils.logger' not in sys.modules:
mock_log = types.ModuleType('app.utils.logger')
mock_log.logger = MagicMock()
sys.modules['app.utils.logger'] = mock_log
module_name = 'app.services.bitget_trading_api_sdk_bal_test'
spec = importlib.util.spec_from_file_location(module_name, sdk_path)
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
return module.BitgetTradingAPI
# ──────────────────────────────────────────────
# 测试 1: fetch_balance 成功,直接返回
# ──────────────────────────────────────────────
def test_fetch_balance_success():
"""fetch_balance 返回 USDT不走原始 API"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {
'USDT': {'free': 1975.0, 'used': 50.0, 'total': 2025.0},
'free': {'USDT': 1975.0},
'used': {'USDT': 50.0},
'total': {'USDT': 2025.0},
}
balance = api.get_balance()
assert balance == {
'USDT': {
'available': '1975.0',
'frozen': '50.0',
'locked': '50.0',
'equity': '2025.0',
}
}
# 不应调原始 API
api.exchange.privateUtaGetV3AccountAssets.assert_not_called()
# ──────────────────────────────────────────────
# 测试 2: fetch_balance 成功,余额为 0
# ──────────────────────────────────────────────
def test_fetch_balance_zero_balance():
"""fetch_balance 返回 0 余额,应正确解析"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {
'USDT': {'free': 0, 'used': 0, 'total': 0},
}
balance = api.get_balance()
assert balance['USDT']['available'] == '0'
assert balance['USDT']['equity'] == '0'
# ──────────────────────────────────────────────
# 测试 3: fetch_balance 无 USDT → 回退 UTA
# ──────────────────────────────────────────────
def test_fetch_balance_no_usdt_fallback_uta():
"""fetch_balance 返回无 USDT → 回退到 UTA 原始 API"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
# fetch_balance 没有USDT
api.exchange.fetch_balance.return_value = {
'BTC': {'free': 0.5, 'used': 0, 'total': 0.5},
}
# UTA 回退返回正常数据
api.exchange.privateUtaGetV3AccountAssets.return_value = {
'code': '00000',
'data': {
'assets': [
{
'coin': 'USDT',
'available': '1975.00',
'locked': '0',
'equity': '1975.00',
}
],
},
}
balance = api.get_balance()
assert balance['USDT']['available'] == '1975.00'
assert balance['USDT']['equity'] == '1975.00'
api.exchange.privateUtaGetV3AccountAssets.assert_called_once()
# ──────────────────────────────────────────────
# 测试 4: fetch_balance 无 USDT → 回退合约 API
# ──────────────────────────────────────────────
def test_fetch_balance_no_usdt_fallback_contract():
"""use_unified_account=False → 回退到合约 API"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = False
# fetch_balance 没有USDT
api.exchange.fetch_balance.return_value = {}
# 合约 API 回退
api.exchange.privateMixGetV2MixAccountAccounts.return_value = {
'code': '00000',
'data': [
{
'marginCoin': 'USDT',
'available': '800.50',
'locked': '100.00',
'accountEquity': '900.50',
}
],
}
balance = api.get_balance()
assert balance['USDT']['available'] == '800.50'
assert balance['USDT']['frozen'] == '100.00'
assert balance['USDT']['equity'] == '900.50'
# ──────────────────────────────────────────────
# 测试 5: fetch_balance 抛异常 → 回退 UTA
# ──────────────────────────────────────────────
def test_fetch_balance_exception_fallback():
"""fetch_balance 抛异常 → 回退到原始 API"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.side_effect = Exception("network error")
api.exchange.privateUtaGetV3AccountAssets.return_value = {
'code': '00000',
'data': {
'assets': [
{
'coin': 'USDT',
'available': '500.00',
'locked': '0',
'equity': '500.00',
}
],
},
}
balance = api.get_balance()
assert balance['USDT']['available'] == '500.00'
api.exchange.privateUtaGetV3AccountAssets.assert_called_once()
# ──────────────────────────────────────────────
# 测试 6: 两个策略都失败 → 返回 {}
# ──────────────────────────────────────────────
def test_both_strategies_fail():
"""fetch_balance 和原始 API 都失败 → 返回空字典"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.side_effect = Exception("timeout")
api.exchange.privateUtaGetV3AccountAssets.side_effect = Exception("auth fail")
balance = api.get_balance()
assert balance == {}
# ──────────────────────────────────────────────
# 测试 7: fetch_balance 返回 free=None
# ──────────────────────────────────────────────
def test_fetch_balance_none_values():
"""fetch_balance 中 free/used 为 None → 解析为 0"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {
'USDT': {'free': None, 'used': None, 'total': None},
}
balance = api.get_balance()
assert balance['USDT']['available'] == '0'
assert balance['USDT']['frozen'] == '0'
assert balance['USDT']['equity'] == '0'
# ──────────────────────────────────────────────
# 测试 8: fetch_balance 返回 free=0, total=0 → equity 用 free+used
# ──────────────────────────────────────────────
def test_fetch_balance_zero_total():
"""total=0 时 equity = free + used"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {
'USDT': {'free': 100, 'used': 50, 'total': 0},
}
balance = api.get_balance()
assert balance['USDT']['equity'] == '150'
# ──────────────────────────────────────────────
# 测试 9: UTA 返回空 assets → 回退到 {}
# ──────────────────────────────────────────────
def test_uta_empty_assets():
"""UTA API 返回空 assets 列表"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {}
api.exchange.privateUtaGetV3AccountAssets.return_value = {
'code': '00000',
'data': {'assets': []},
}
balance = api.get_balance()
assert balance == {}
# ──────────────────────────────────────────────
# 测试 10: get_account_state 端到端解析
# ──────────────────────────────────────────────
def test_get_account_state_e2e():
"""get_account_state 正确解析 balance 返回值"""
BitgetTradingAPI = load_bitget_sdk_class()
# Mock get_bitget_live_service 的依赖链
# 直接测 get_account_state 解析逻辑
from app.services.bitget_live_trading_service import CONTRACT_SIZES
# 模拟 balance 返回
balance_return = {
'USDT': {
'available': '1975.50',
'frozen': '24.50',
'locked': '24.50',
'equity': '2000.00',
}
}
# 直接验证解析逻辑(不初始化完整 service
usdt = balance_return.get('USDT', {})
available = float(usdt.get('available', 0) or 0)
frozen = float(usdt.get('frozen', 0) or 0)
equity = float(usdt.get('equity', 0) or 0)
account_value = equity if equity > 0 else (available + frozen)
assert available == 1975.50
assert frozen == 24.50
assert equity == 2000.00
assert account_value == 2000.00
# ──────────────────────────────────────────────
# 测试 11: balance 返回 {} 时 get_account_state 各字段为 0
# ──────────────────────────────────────────────
def test_get_account_state_empty_balance():
"""balance 返回空 → 全部为 0"""
balance_return = {}
usdt = balance_return.get('USDT', {})
available = float(usdt.get('available', 0) or 0)
frozen = float(usdt.get('frozen', 0) or 0)
equity = float(usdt.get('equity', 0) or 0)
account_value = equity if equity > 0 else (available + frozen)
assert available == 0.0
assert frozen == 0.0
assert account_value == 0.0
# ──────────────────────────────────────────────
# 测试 12: fetch_balance 返回字符串数值
# ──────────────────────────────────────────────
def test_fetch_balance_string_values():
"""ccxt 有时返回字符串类型的余额"""
BitgetTradingAPI = load_bitget_sdk_class()
api = BitgetTradingAPI.__new__(BitgetTradingAPI)
api.exchange = MagicMock()
api.use_unified_account = True
api.exchange.fetch_balance.return_value = {
'USDT': {'free': '1975.50', 'used': '24.50', 'total': '2000.00'},
}
balance = api.get_balance()
assert balance['USDT']['available'] == '1975.50'
assert balance['USDT']['equity'] == '2000.00'