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

1# brokers/korea_investment/korea_invest_env.py 

2 

3import requests 

4import json 

5import os 

6from datetime import datetime, timedelta 

7import certifi 

8import yaml 

9import logging 

10import pytz 

11from brokers.korea_investment.korea_invest_token_provider import TokenProvider 

12 

13 

14class KoreaInvestApiEnv: 

15 """ 

16 한국투자증권 Open API 환경 설정을 관리하는 클래스입니다. 

17 API 키, 계좌 정보, 도메인 정보 등을 로드하고, 

18 API 요청에 필요한 기본 헤더를 생성합니다. 

19 토큰을 로컬 파일에 저장하고 재사용하는 기능을 포함합니다. 

20 """ 

21 

22 def __init__(self, config_data, logger=None): 

23 self._config_data = config_data 

24 self._logger = logger if logger else logging.getLogger(__name__) 

25 

26 self._token_provider = None 

27 self._token_provider_real = TokenProvider(token_file_path="config/token_real.json", logger=self._logger) 

28 self._token_provider_paper = TokenProvider(token_file_path="config/token_paper.json", logger=self._logger) 

29 self._load_config() 

30 self._set_base_urls() # 초기 base_url, websocket_url 설정 (config.yaml의 is_paper_trading 기반) 

31 self._token_file_path = os.path.join(os.getcwd(), 'kis_access_token.yaml') 

32 self._session = requests.Session() 

33 

34 self.is_paper_trading = None 

35 self._base_url = None 

36 self._websocket_url = None 

37 self.active_config = None 

38 

39 def _load_config(self): 

40 self.api_key = self._config_data.get('api_key') 

41 self.api_secret_key = self._config_data.get('api_secret_key') 

42 self.stock_account_number = self._config_data.get('stock_account_number') 

43 

44 self.paper_api_key = self._config_data.get('paper_api_key') 

45 self.paper_api_secret_key = self._config_data.get('paper_api_secret_key') 

46 self.paper_stock_account_number = self._config_data.get('paper_stock_account_number') 

47 

48 self.htsid = self._config_data.get('htsid') 

49 self.custtype = self._config_data.get('custtype', 'P') 

50 

51 self.is_paper_trading = self._config_data.get('is_paper_trading', False) 

52 

53 self.my_agent = self._config_data.get('my_agent', 

54 "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36") 

55 

56 def _set_base_urls(self): 

57 """is_paper_trading 값에 따라 base_url과 websocket_url을 설정합니다.""" 

58 if self.is_paper_trading: 

59 self._base_url = self._config_data.get('paper_url') 

60 self._websocket_url = self._config_data.get('paper_websocket_url') 

61 else: 

62 self._base_url = self._config_data.get('url') 

63 self._websocket_url = self._config_data.get('websocket_url') 

64 

65 if not self._base_url or not self._websocket_url: 

66 raise ValueError("API URL 또는 WebSocket URL이 config.yaml에 올바르게 설정되지 않았습니다.") 

67 

68 def set_trading_mode(self, is_paper: bool): 

69 """ 

70 거래 모드(실전/모의)를 동적으로 변경합니다. 

71 :param is_paper: 모의투자 모드이면 True, 실전 투자 모드이면 False 

72 """ 

73 # if self._is_paper_trading != is_paper: 

74 if self.is_paper_trading is not is_paper: 

75 self.is_paper_trading = is_paper 

76 self._set_base_urls() 

77 self.active_config = self.get_full_config() 

78 

79 if self.is_paper_trading is True: 

80 self._token_provider = self._token_provider_paper 

81 else: 

82 self._token_provider = self._token_provider_real 

83 self._logger.info(f"거래 모드가 {'모의투자' if is_paper else '실전투자'} 환경으로 변경되었습니다.") 

84 else: 

85 self._logger.info(f"거래 모드가 이미 {'모의투자' if is_paper else '실전투자'} 환경으로 설정되어 있습니다.") 

86 

87 def get_base_headers(self): 

88 """API 요청 시 사용할 기본 헤더를 반환합니다.""" 

89 headers = { 

90 "Content-Type": "application/json", 

91 "User-Agent": self.my_agent, 

92 "charset": "UTF-8" 

93 } 

94 return headers 

95 

96 def get_full_config(self): 

97 """ 

98 현재 활성화된 환경(실전/모의투자)에 맞는 API 키, 계좌 정보, URL 등을 반환합니다. 

99 tr_ids는 config_data에서 그대로 가져와 포함합니다. 

100 """ 

101 active_api_key = self.paper_api_key if self.is_paper_trading else self.api_key 

102 active_api_secret_key = self.paper_api_secret_key if self.is_paper_trading else self.api_secret_key 

103 active_stock_account_number = self.paper_stock_account_number if self.is_paper_trading else self.stock_account_number 

104 active_base_url = self._base_url 

105 active_websocket_url = self._websocket_url 

106 

107 tr_ids_from_config = self._config_data.get('tr_ids', {}) 

108 

109 # TokenProvider 에서 현재 활성 토큰과 만료 시간을 가져옴 

110 # current_access_token = self._token_provider._access_token # 직접 접근 

111 # current_token_expired_at = self._token_provider._token_expired_at # 직접 접근 

112 

113 return { 

114 'api_key': active_api_key, 

115 'api_secret_key': active_api_secret_key, 

116 'stock_account_number': active_stock_account_number, 

117 'base_url': active_base_url, 

118 'websocket_url': active_websocket_url, 

119 'htsid': self.htsid, 

120 'custtype': self.custtype, 

121 # 'access_token': current_access_token, # 현재 인스턴스에 저장된 토큰 

122 # 'token_expired_at': current_token_expired_at, # 현재 인스턴스에 저장된 만료 시간 

123 'is_paper_trading': self.is_paper_trading, 

124 'tr_ids': tr_ids_from_config, 

125 'paths': self._config_data['paths'], 

126 'params': self._config_data['params'], 

127 # 'market_code': self._config_data['market_code'], 

128 '_env_instance': self # <--- _env_instance는 KoreaInvestAPI로 전달하기 위해 여기에 추가 

129 } 

130 

131 async def get_access_token(self, force_new=False): 

132 """ 

133 접근 토큰을 발급받거나 갱신합니다. 

134 토큰 관리를 TokenProvider에 위임합니다. 

135 """ 

136 # TokenProvider의 get_access_token을 호출하고 결과를 반환합니다. 

137 # TokenProvider 내부에서 force_new 로직을 처리하므로, 여기서는 인자를 전달하지 않습니다. 

138 if not self._token_provider: 

139 raise RuntimeError("TokenProvider가 초기화되지 않았습니다. set_trading_mode 먼저 호출하세요.") 

140 

141 return await self._token_provider.get_access_token( 

142 base_url=self.active_config['base_url'], 

143 app_key=self.active_config['api_key'], 

144 app_secret=self.active_config['api_secret_key'] 

145 ) 

146 

147 def save_access_token(self, token: str): 

148 if not self._token_provider: 

149 raise RuntimeError("TokenProvider가 초기화되지 않았습니다.") 

150 self._token_provider.save_access_token(token) 

151 

152 def invalidate_token(self): 

153 if self._token_provider: 

154 self._token_provider.invalidate_token() 

155 

156 async def refresh_token(self): 

157 if self._token_provider: 

158 await self._token_provider.refresh_token( 

159 base_url=self.active_config['base_url'], 

160 app_key=self.active_config['api_key'], 

161 app_secret=self.active_config['api_secret_key'] 

162 ) 

163 

164 def get_base_url(self): 

165 return self._base_url 

166 

167 def get_real_base_url(self) -> str: 

168 """항상 실전 base URL 반환 (조회 API용).""" 

169 return self._config_data.get('url') 

170 

171 def get_real_config(self) -> dict: 

172 """항상 실전 API 키/시크릿/URL 반환.""" 

173 return { 

174 'api_key': self.api_key, 

175 'api_secret_key': self.api_secret_key, 

176 'base_url': self._config_data.get('url'), 

177 } 

178 

179 async def get_real_access_token(self) -> str: 

180 """항상 실전 토큰 반환 (_token_provider_real 사용).""" 

181 real_cfg = self.get_real_config() 

182 return await self._token_provider_real.get_access_token( 

183 base_url=real_cfg['base_url'], 

184 app_key=real_cfg['api_key'], 

185 app_secret=real_cfg['api_secret_key'] 

186 ) 

187 

188 def get_websocket_url(self): 

189 return self._websocket_url