테스트 강화 및 코드 품질 개선

verify_krw_budget.py

@@ -0,0 +1,123 @@
+"""
+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)