Coverage for django_napse/core/models/fleets/cluster.py: 84%

1from django.db import models 

2 

3# from django_napse.core.models.connections.connection import Connection 

4from django_napse.core.models.fleets.link import Link 

5from django_napse.core.models.transactions.transaction import Transaction 

6from django_napse.utils.constants import TRANSACTION_TYPES 

7from django_napse.utils.errors import BotError, ClusterError 

8 

9 

10class Cluster(models.Model): 

11 fleet = models.ForeignKey( 

12 "Fleet", 

13 on_delete=models.CASCADE, 

14 related_name="clusters", 

15 ) 

16 template_bot = models.OneToOneField( 

17 "Bot", 

18 on_delete=models.CASCADE, 

19 related_name="cluster", 

20 ) 

21 share = models.FloatField() 

22 breakpoint = models.FloatField() 

23 autoscale = models.BooleanField() 

24 

25 def __str__(self): 

26 return f"Cluster: {self.fleet}" 

27 

28 def save(self, *args, **kwargs): 

29 if not self.config.immutable: 

30 error_msg = "In a fleet, the config must be immutable." 

31 raise ClusterError.MutableBotConfig(error_msg) 

32 return super().save(*args, **kwargs) 

33 

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

35 string = "" 

36 string += f"{beacon}Cluster {self.pk}:\n" 

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

38 string += f"{beacon}\t{self.fleet=}\n" 

39 string += f"{beacon}\t{self.template_bot=}\n" 

40 string += f"{beacon}\t{self.share=}\n" 

41 string += f"{beacon}\t{self.breakpoint=}\n" 

42 string += f"{beacon}\t{self.autoscale=}\n" 

43 

44 new_beacon = beacon + "\t" 

45 string += f"{beacon}Links:\n" 

46 for link in self.links.all(): 

47 string += f"{link.info(verbose=False, beacon=new_beacon)}\n" 

48 

49 string += f"{beacon}Connections:\n" 

50 connections = [] 

51 for link in self.links.all(): 

52 connections += list(link.bot.connections.all()) 

53 for connection in connections: 

54 string += f"{connection.info(verbose=False, beacon=new_beacon)}\n" 

55 

56 if verbose: # pragma: no cover 

57 print(string) 

58 return string 

59 

60 @property 

61 def config(self): 

62 return self.template_bot.strategy.config 

63 

64 def invest(self, space, amount, ticker): 

65 all_connections = [] 

66 bots = [link.bot for link in self.links.all().order_by("importance")] 

67 if len(bots) == 0: 

68 new_bot = self.template_bot.copy() 

69 Link.objects.create(bot=new_bot, cluster=self, importance=1) 

70 # connection = Connection.objects.create(bot=new_bot, owner=space.wallet) 

71 connection = space.wallet.connect_to_bot(new_bot) 

72 Transaction.objects.create( 

73 from_wallet=space.wallet, 

74 to_wallet=connection.wallet, 

75 amount=amount, 

76 ticker=ticker, 

77 transaction_type=TRANSACTION_TYPES.CONNECTION_DEPOSIT, 

78 ) 

79 all_connections.append(connection) 

80 elif len(bots) == 1: 

81 bot = bots[0] 

82 if ticker not in bot.architecture.accepted_investment_tickers() or ticker not in bot.architecture.accepted_tickers(): 

83 error_msg = f"Bot {bot} does not accept ticker {ticker}." 

84 raise BotError.InvalidTicker(error_msg) 

85 

86 connection = space.wallet.connect_to_bot(bot) 

87 # connection = Connection.objects.get(bot=bot, owner=space.wallet) 

88 Transaction.objects.create( 

89 from_wallet=space.wallet, 

90 to_wallet=connection.wallet, 

91 amount=amount, 

92 ticker=ticker, 

93 transaction_type=TRANSACTION_TYPES.CONNECTION_DEPOSIT, 

94 ) 

95 all_connections.append(connection) 

96 else: 

97 error_msg = "Autoscale not implemented yet." 

98 raise NotImplementedError(error_msg) 

99 return all_connections