Coverage for brokers / korea_investment / korea_invest_account_api.py: 100%

1# brokers/korea_investment/korea_invest_account_api.py 

2 

3import httpx 

4from brokers.korea_investment.korea_invest_api_base import KoreaInvestApiBase 

5from brokers.korea_investment.korea_invest_env import KoreaInvestApiEnv 

6from brokers.korea_investment.korea_invest_params_provider import Params 

7from brokers.korea_investment.korea_invest_header_provider import KoreaInvestHeaderProvider 

8from brokers.korea_investment.korea_invest_url_provider import KoreaInvestUrlProvider 

9from brokers.korea_investment.korea_invest_url_keys import EndpointKey 

10from brokers.korea_investment.korea_invest_trid_provider import KoreaInvestTrIdProvider 

11from typing import Optional 

12from common.types import ResCommonResponse, Exchange 

13 

14 

15class KoreaInvestApiAccount(KoreaInvestApiBase): 

16 def __init__(self, 

17 env: KoreaInvestApiEnv, 

18 logger=None, 

19 market_clock=None, 

20 async_client: Optional[httpx.AsyncClient] = None, 

21 header_provider: Optional[KoreaInvestHeaderProvider] = None, 

22 url_provider: Optional[KoreaInvestUrlProvider] = None, 

23 trid_provider: Optional[KoreaInvestTrIdProvider] = None): 

24 super().__init__(env, 

25 logger, 

26 market_clock, 

27 async_client=async_client, 

28 header_provider=header_provider, 

29 url_provider=url_provider, 

30 trid_provider=trid_provider) 

31 

32 async def get_account_balance(self, exchange: Exchange = Exchange.KRX) -> ResCommonResponse: 

33 """ 

34 모의투자 또는 실전투자 계좌의 잔고를 조회하는 메서드. 

35 투자환경(env)에 따라 자동 분기된다. 

36 exchange=Exchange.NXT 인 경우 시간외단일가 파라미터(AFHR_FLPR_YN)를 "X"로 설정한다. 

37 """ 

38 full_config = self._env.active_config 

39 

40 is_paper = self._env.is_paper_trading # 모의투자 여부 판단 

41 tr_id = self._trid_provider.account_inquire_balance() # 모드에 따라 자동 

42 

43 # ✅ 요청 단위 임시 헤더 주입 

44 self._headers.set_tr_id(tr_id) 

45 self._headers.set_custtype(full_config['custtype']) 

46 

47 full_account_number = full_config['stock_account_number'] 

48 if '-' in full_account_number and len(full_account_number.split('-')[1]) == 2: 

49 cano, acnt_prdt_cd = full_account_number.split('-') 

50 else: 

51 cano = full_account_number 

52 acnt_prdt_cd = "01" 

53 

54 # NXT 거래소는 AFHR_FLPR_YN을 "X"로 설정 

55 afhr_flpr_yn = "X" if exchange == Exchange.NXT else "N" 

56 params = Params.account_balance(cano=cano, acnt_prdt_cd=acnt_prdt_cd, afhr_flpr_yn=afhr_flpr_yn) 

57 

58 mode_str = "모의투자" if is_paper else "실전투자" 

59 self._logger.info(f"{mode_str} 계좌 잔고 조회 시도...") 

60 return await self.call_api('GET', EndpointKey.INQUIRE_BALANCE, params=params, retry_count=3)