Coverage for django_napse/core/models/wallets/wallet.py: 95%

1import time 

2 

3from django.db import models 

4 

5from django_napse.core.models.bots.controller import Controller 

6from django_napse.core.models.connections.connection import Connection 

7from django_napse.core.models.wallets.currency import Currency 

8from django_napse.core.models.wallets.managers import WalletManager 

9from django_napse.utils.errors import WalletError 

10from django_napse.utils.findable_class import FindableClass 

11 

12 

13class Wallet(models.Model, FindableClass): 

14 title = models.CharField(max_length=255, default="Wallet") 

15 locked = models.BooleanField(default=False) 

16 created_at = models.DateTimeField(auto_now_add=True) 

17 

18 objects = WalletManager() 

19 

20 def __str__(self): 

21 return f"WALLET: {self.pk=}" 

22 

23 def info(self, verbose=True, beacon=""): 

24 self = self.find() 

25 string = "" 

26 string += f"{beacon}Wallet ({self.pk=}):\t{type(self)}\n" 

27 string += f"{beacon}Args:\n" 

28 string += f"{beacon}\t{self.title=}\n" 

29 string += f"{beacon}\t{self.testing=}\n" 

30 string += f"{beacon}\t{self.locked=}\n" 

31 string += f"{beacon}Currencies\n" 

32 if self.currencies.count() == 0: 

33 string += f"{beacon}\tNo currencies\n" 

34 else: 

35 for currency in self.currencies.all(): 

36 string += f"{beacon}\t{currency.ticker}: {currency.amount} @ ${currency.mbp}\n" 

37 if verbose: # pragma: no cover 

38 print(string) 

39 return string 

40 

41 @property 

42 def testing(self): 

43 return self.owner.testing 

44 

45 @property 

46 def space(self): # pragma: no cover 

47 error_msg = f"space() not implemented by default. Please implement it in {self.__class__}." 

48 raise NotImplementedError(error_msg) 

49 

50 @property 

51 def exchange_account(self): # pragma: no cover 

52 error_msg = "exchange_account() not implemented by default. Please implement in a subclass of Wallet." 

53 raise NotImplementedError(error_msg) 

54 

55 def spend(self, amount: float, ticker: str, recv: int = 3, **kwargs) -> None: 

56 if not kwargs.get("force", False): 

57 error_msg = "DANGEROUS: You should not use this method outside of select circumstances. Use Transactions instead." 

58 raise WalletError.SpendError(error_msg) 

59 

60 if amount <= 0: 

61 error_msg: str = f"Amount must be positive, got {amount}" 

62 raise ValueError(error_msg) 

63 

64 start_time = time.time() 

65 while self.locked: 

66 time.sleep(0.01) 

67 if time.time() - start_time > recv: 

68 error_msg: str = f"Wallet: {self.title} is locked" 

69 raise TimeoutError(error_msg) 

70 self.locked = True 

71 self.save() 

72 

73 try: 

74 currency = self.currencies.get(ticker=ticker) 

75 except Currency.DoesNotExist as err: 

76 self.locked = False 

77 self.save() 

78 error_msg: str = f"Currency {ticker} does not exist in wallet: {self.title}." 

79 raise WalletError.SpendError(error_msg) from err 

80 

81 if currency.amount < amount: 

82 self.locked = False 

83 self.save() 

84 error_msg: str = f"Not enough money in wallet: {self.title} ({amount} > {currency.amount} {ticker})." 

85 raise WalletError.SpendError(error_msg) 

86 

87 currency.amount -= amount 

88 currency.save() 

89 self.locked = False 

90 self.save() 

91 

92 def top_up(self, amount: float, ticker: str, mbp: float | None = None, recv: int = 3, **kwargs) -> None: 

93 if not kwargs.get("force", False): 

94 error_msg = "DANGEROUS: You should not use this method outside of select circumstances. Use Transactions instead." 

95 raise WalletError.TopUpError(error_msg) 

96 

97 if amount <= 0: 

98 error_msg: str = f"Amount must be positive, got {amount}" 

99 raise ValueError(error_msg) 

100 

101 if mbp is None: 

102 mbp = Controller.get_asset_price(exchange_account=self.exchange_account, base=ticker) 

103 

104 start_time = time.time() 

105 while self.locked: 

106 time.sleep(0.01) 

107 if time.time() - start_time > recv: 

108 error_msg: str = f"Wallet: {self.title} is locked" 

109 raise TimeoutError(error_msg) 

110 

111 self.locked = True 

112 self.save() 

113 try: 

114 currency = self.currencies.get(ticker=ticker) 

115 except Currency.DoesNotExist: 

116 currency = Currency.objects.create(wallet=self, ticker=ticker, amount=0, mbp=0) 

117 currency.mbp = (currency.mbp * currency.amount + mbp * amount) / (currency.amount + amount) 

118 currency.amount += amount 

119 currency.save() 

120 self.locked = False 

121 self.save() 

122 

123 def has_funds(self, amount: float, ticker: str) -> bool: 

124 try: 

125 currency = self.currencies.get(ticker=ticker) 

126 except Currency.DoesNotExist: 

127 return False 

128 return currency.amount >= amount 

129 

130 def get_amount(self, ticker: str) -> float: 

131 try: 

132 curr = self.currencies.get(ticker=ticker) 

133 except Currency.DoesNotExist: 

134 return 0 

135 return curr.amount 

136 

137 def value_mbp(self) -> float: 

138 value = 0 

139 for currency in self.currencies.all(): 

140 if currency.amount == 0: 

141 continue 

142 value += currency.amount * currency.mbp 

143 return value 

144 

145 def value_market(self) -> float: 

146 value = 0 

147 for currency in self.currencies.all(): 

148 if currency.amount == 0: 

149 continue 

150 value += currency.amount * Controller.get_asset_price(exchange_account=self.exchange_account, base=currency.ticker) 

151 return value 

152 

153 def to_dict(self): 

154 currencies = self.currencies.all() 

155 return { 

156 "title": self.title, 

157 "testing": self.testing, 

158 "locked": self.locked, 

159 "created_at": self.created_at, 

160 "currencies": { 

161 currency.ticker: { 

162 "amount": currency.amount, 

163 "mbp": currency.mbp, 

164 } 

165 for currency in currencies 

166 }, 

167 } 

168 

169 

170class SpaceWallet(Wallet): 

171 owner = models.OneToOneField("NapseSpace", on_delete=models.CASCADE, related_name="wallet") 

172 

173 def __str__(self): 

174 return f"WALLET: {self.pk=}\nOWNER: {self.owner=}" 

175 

176 @property 

177 def space(self): 

178 return self.owner 

179 

180 @property 

181 def exchange_account(self): 

182 return self.space.exchange_account.find() 

183 

184 def connect_to_bot(self, bot): 

185 try: 

186 connection = self.connections.get(owner=self, bot=bot) 

187 except Connection.DoesNotExist: 

188 connection = Connection.objects.create(owner=self, bot=bot) 

189 return connection 

190 

191 

192class SpaceSimulationWallet(Wallet): 

193 owner = models.OneToOneField("NapseSpace", on_delete=models.CASCADE, related_name="simulation_wallet") 

194 

195 def __str__(self): 

196 return f"WALLET: {self.pk=}\nOWNER: {self.owner=}" 

197 

198 @property 

199 def testing(self): 

200 return True 

201 

202 @property 

203 def space(self): 

204 return self.owner 

205 

206 @property 

207 def exchange_account(self): 

208 return self.space.exchange_account.find() 

209 

210 def reset(self): 

211 self.currencies.all().delete() 

212 

213 def connect_to_bot(self, bot): 

214 """Get or create connection to bot.""" 

215 try: 

216 connection = self.connections.get(owner=self, bot=bot) 

217 except Connection.DoesNotExist: 

218 connection = Connection.objects.create(owner=self, bot=bot) 

219 return connection 

220 

221 

222class OrderWallet(Wallet): 

223 owner = models.OneToOneField("Order", on_delete=models.CASCADE, related_name="wallet") 

224 

225 def __str__(self): 

226 return f"WALLET: {self.pk=}\nOWNER: {self.owner=}" 

227 

228 @property 

229 def exchange_account(self): 

230 return self.owner.exchange_account.find() 

231 

232 

233class ConnectionWallet(Wallet): 

234 owner = models.OneToOneField("Connection", on_delete=models.CASCADE, related_name="wallet") 

235 

236 def __str__(self): 

237 return f"WALLET: {self.pk=}\nOWNER: {self.owner=}" 

238 

239 @property 

240 def space(self): 

241 return self.owner.space 

242 

243 @property 

244 def exchange_account(self): 

245 return self.space.exchange_account.find()