Coverage for brokers / korea_investment / korea_invest_header_provider.py: 97%

1# brokers/korea_investment/korea_invest_header_provider.py 

2from __future__ import annotations 

3from dataclasses import dataclass, field 

4from typing import Dict, Optional, Iterator 

5from contextlib import contextmanager 

6import os 

7 

8 

9@dataclass 

10class KoreaInvestHeaderProvider: 

11 """ 

12 중앙집중 헤더 관리자. 

13 - 기본 헤더 보유 

14 - 토큰/앱키/앱시크릿 자동 주입 

15 - TR-ID, custtype, hashkey, gt_uid 등 단건 API 실행 전/중 임시 설정(Context) 지원 

16 - dict를 직접 노출하지 않고, 매 호출 시 `build()`로 사본 생성 

17 """ 

18 my_agent: str 

19 appkey: str = "" 

20 appsecret: str = "" 

21 custtype_default: str = "P" 

22 

23 # 내부 상태(임시/지속) 분리 

24 _base: Dict[str, str] = field(default_factory=dict) 

25 _volatile: Dict[str, str] = field(default_factory=dict) 

26 

27 def __post_init__(self): 

28 self._base = { 

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

30 "User-Agent": self.my_agent, 

31 "charset": "UTF-8", 

32 "Authorization": "", 

33 "appkey": self.appkey, 

34 "appsecret": self.appsecret, 

35 # custtype은 보통 공통이지만, 엔드포인트별 변경 가능 → volatile에 넣는 편의 메서드 제공 

36 } 

37 self._volatile = {"custtype": self.custtype_default} 

38 

39 # --------- 지속 필드 설정 --------- 

40 def set_auth_bearer(self, access_token: str) -> None: 

41 self._base["Authorization"] = f"Bearer {access_token}" 

42 

43 def set_app_keys(self, appkey: str, appsecret: str) -> None: 

44 self._base["appkey"] = appkey 

45 self._base["appsecret"] = appsecret 

46 

47 # --------- 임시(요청별) 필드 설정 --------- 

48 def set_tr_id(self, tr_id: Optional[str]) -> None: 

49 if tr_id is None: 

50 self._volatile.pop("tr_id", None) 

51 else: 

52 self._volatile["tr_id"] = tr_id 

53 

54 def set_custtype(self, custtype: Optional[str]) -> None: 

55 if custtype is None: 

56 self._volatile.pop("custtype", None) 

57 else: 

58 self._volatile["custtype"] = custtype 

59 

60 def set_hashkey(self, hashkey: Optional[str]) -> None: 

61 if hashkey is None: 

62 self._volatile.pop("hashkey", None) 

63 else: 

64 self._volatile["hashkey"] = hashkey 

65 

66 def set_gt_uid(self, gt_uid: Optional[str] = None) -> None: 

67 # 주지 않을 경우 자동 생성 

68 if gt_uid is None: 

69 gt_uid = os.urandom(16).hex() 

70 self._volatile["gt_uid"] = gt_uid 

71 

72 def clear_order_headers(self) -> None: 

73 # 주문계열에서 쓰는 hashkey/gt_uid 등만 정리 

74 for k in ("hashkey", "gt_uid"): 

75 self._volatile.pop(k, None) 

76 

77 def sync_from_env(self, env) -> None: 

78 """모드가 정해진 이후 env.active_config로부터 키/모드 동기화.""" 

79 cfg = getattr(env, "active_config", None) or {} 

80 self.set_app_keys(cfg.get("api_key", ""), cfg.get("api_secret_key", "")) 

81 self.set_custtype(cfg.get("custtype", "P")) 

82 

83 # --------- 빌드/컨텍스트 --------- 

84 def build(self) -> Dict[str, str]: 

85 # 매 요청 시 사본 생성(외부 변조 방지) 

86 h = {**self._base, **self._volatile} 

87 # 값이 빈 문자열인 키는 제거(헤더 깔끔하게) 

88 return {k: v for k, v in h.items() if v is not None and v != ""} 

89 

90 @contextmanager 

91 def temp(self, *, tr_id: Optional[str] = None, custtype: Optional[str] = None, 

92 hashkey: Optional[str] = None, gt_uid: Optional[str] = None) -> Iterator[None]: 

93 """요청 단위 임시 헤더 설정 컨텍스트. 

94 사용 예) 

95 with header_mgr.temp(tr_id="FHK...", custtype="P"): 

96 client.get(..., headers=header_mgr.build()) 

97 컨텍스트 종료 시 원복. 

98 """ 

99 backup = dict(self._volatile) 

100 try: 

101 if tr_id is not None: 101 ↛ 103line 101 didn't jump to line 103 because the condition on line 101 was always true

102 self.set_tr_id(tr_id) 

103 if custtype is not None: 

104 self.set_custtype(custtype) 

105 if hashkey is not None: 

106 self.set_hashkey(hashkey) 

107 if gt_uid is not None: 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true

108 self.set_gt_uid(gt_uid) 

109 yield 

110 finally: 

111 self._volatile = backup 

112 

113 def fork(self) -> "KoreaInvestHeaderProvider": 

114 cloned = KoreaInvestHeaderProvider( 

115 my_agent=self._base.get("User-Agent", self.my_agent), 

116 appkey=self._base.get("appkey", ""), 

117 appsecret=self._base.get("appsecret", ""), 

118 custtype_default=self._volatile.get("custtype", self.custtype_default), 

119 ) 

120 cloned._base = dict(self._base) 

121 cloned._volatile = dict(self._volatile) 

122 return cloned 

123 

124# ----------------------------------------------------------------------------- 

125# 사용 편의 팩토리 

126# ----------------------------------------------------------------------------- 

127def build_header_provider_from_env(env) -> KoreaInvestHeaderProvider: 

128 # 생성 시엔 UA만. 키/모드는 나중에 sync_from_env로 주입 

129 return KoreaInvestHeaderProvider( 

130 my_agent=getattr(env, "my_agent", "python-client") 

131 )