124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
"""
|
|
KRWBudgetManager 동작 검증 스크립트
|
|
"""
|
|
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
from src.common import KRWBudgetManager
|
|
|
|
|
|
def test_basic():
|
|
"""기본 동작 테스트"""
|
|
print("\n=== 기본 동작 테스트 ===")
|
|
manager = KRWBudgetManager()
|
|
|
|
# Mock Upbit 객체
|
|
class MockUpbit:
|
|
def get_balance(self, currency):
|
|
return 100000
|
|
|
|
upbit = MockUpbit()
|
|
|
|
# 테스트 1: 전액 할당
|
|
success, allocated = manager.allocate("KRW-BTC", 50000, upbit)
|
|
print(f"테스트 1 - 전액 할당: success={success}, allocated={allocated}")
|
|
assert success and allocated == 50000, "전액 할당 실패"
|
|
|
|
# 테스트 2: 부분 할당 (남은 50000원)
|
|
success, allocated = manager.allocate("KRW-ETH", 60000, upbit)
|
|
print(f"테스트 2 - 부분 할당: success={success}, allocated={allocated}")
|
|
assert success and allocated == 50000, "부분 할당 실패"
|
|
|
|
# 테스트 3: 할당 실패 (잔고 없음)
|
|
success, allocated = manager.allocate("KRW-XRP", 10000, upbit)
|
|
print(f"테스트 3 - 할당 실패: success={success}, allocated={allocated}")
|
|
assert not success and allocated == 0, "할당 실패 처리 오류"
|
|
|
|
print("✅ 기본 동작 테스트 통과\n")
|
|
manager.clear()
|
|
|
|
|
|
def test_concurrency():
|
|
"""동시성 테스트"""
|
|
print("=== 동시성 테스트 ===")
|
|
manager = KRWBudgetManager()
|
|
|
|
class MockUpbit:
|
|
def get_balance(self, currency):
|
|
return 100000
|
|
|
|
upbit = MockUpbit()
|
|
results = []
|
|
|
|
def worker(symbol, amount):
|
|
success, allocated = manager.allocate(symbol, amount, upbit)
|
|
results.append((symbol, success, allocated))
|
|
if success:
|
|
time.sleep(0.01) # 주문 시뮬레이션
|
|
manager.release(symbol)
|
|
|
|
# 3개 스레드가 동시에 50000원씩 요청 (총 150000 > 100000)
|
|
threads = [threading.Thread(target=worker, args=(f"KRW-COIN{i}", 50000)) for i in range(3)]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# 검증
|
|
total_allocated = sum(allocated for _, success, allocated in results if success)
|
|
print(f"총 요청: 150000원, 총 할당: {total_allocated}원")
|
|
print(f"결과: {results}")
|
|
|
|
assert total_allocated <= 100000, f"과할당 발생: {total_allocated}"
|
|
print("✅ 동시성 테스트 통과\n")
|
|
manager.clear()
|
|
|
|
|
|
def test_release():
|
|
"""예산 해제 테스트"""
|
|
print("=== 예산 해제 테스트 ===")
|
|
manager = KRWBudgetManager()
|
|
|
|
class MockUpbit:
|
|
def get_balance(self, currency):
|
|
return 100000
|
|
|
|
upbit = MockUpbit()
|
|
|
|
# 할당
|
|
success1, _ = manager.allocate("KRW-BTC", 50000, upbit)
|
|
print(f"BTC 할당: {manager.get_allocations()}")
|
|
|
|
# 해제
|
|
manager.release("KRW-BTC")
|
|
print(f"BTC 해제 후: {manager.get_allocations()}")
|
|
|
|
# 재할당 가능
|
|
success2, allocated = manager.allocate("KRW-ETH", 50000, upbit)
|
|
print(f"ETH 재할당: success={success2}, allocated={allocated}")
|
|
|
|
assert success1 and success2 and allocated == 50000, "해제 후 재할당 실패"
|
|
print("✅ 예산 해제 테스트 통과\n")
|
|
manager.clear()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
test_basic()
|
|
test_concurrency()
|
|
test_release()
|
|
print("\n🎉 모든 테스트 통과!")
|
|
sys.exit(0)
|
|
except AssertionError as e:
|
|
print(f"\n❌ 테스트 실패: {e}")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"\n❌ 예외 발생: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
sys.exit(1)
|