Coverage for django_napse/simulations/models/simulations/managers/simulation.py: 90%

1from django.apps import apps 

2from django.db import models 

3 

4from django_napse.utils.errors import SimulationError 

5 

6 

7class SimulationManager(models.Manager): 

8 def create(self, space, bot, start_date, end_date, simulation_reference, data): 

9 SimulationDataPoint = apps.get_model("django_napse_simulations", "SimulationDataPoint") 

10 SimulationDataPointExtraInfo = apps.get_model("django_napse_simulations", "SimulationDataPointExtraInfo") 

11 

12 simulation = self.model( 

13 space=space, 

14 bot=bot, 

15 start_date=start_date, 

16 end_date=end_date, 

17 simulation_reference=simulation_reference, 

18 ) 

19 simulation.save() 

20 

21 if data == {}: 

22 return simulation 

23 must_have = ["dates", "values", "actions", "amounts", "tickers"] 

24 for key in must_have: 

25 if key not in data: 

26 error_msg = f"Key {key} not in data" 

27 raise ValueError(error_msg) 

28 

29 first_length = len(data["dates"]) 

30 for key, value in data.items(): 

31 if len(value) != first_length: 

32 error_msg = f"Key {key} has length {len(value)} but should have length {first_length}" 

33 raise SimulationError.InvalidData(error_msg) 

34 

35 extra_info_keys = [key for key in data if key not in must_have] 

36 

37 simulation.save() 

38 bulk_list_data_point = [] 

39 bulk_list_extra_info = [] 

40 for info in zip( 

41 data["dates"], 

42 data["values"], 

43 data["actions"], 

44 data["amounts"], 

45 data["tickers"], 

46 *[data[key] for key in extra_info_keys], 

47 strict=True, 

48 ): 

49 date = info[0] 

50 value = info[1] 

51 action = info[2] 

52 amount = info[3] 

53 ticker = info[4] 

54 extra_info = {key: info[i + 5] for i, key in enumerate(extra_info_keys)} 

55 bulk_list_data_point.append( 

56 SimulationDataPoint( 

57 simulation=simulation, 

58 date=date, 

59 value=value, 

60 action=action, 

61 amount=amount, 

62 ticker=ticker, 

63 ), 

64 ) 

65 for key, value in extra_info.items(): 

66 bulk_list_extra_info.append( 

67 SimulationDataPointExtraInfo( 

68 data_point=bulk_list_data_point[-1], 

69 key=key, 

70 value=str(value), 

71 target_type=type(value).__name__, 

72 ), 

73 ) 

74 

75 if len(bulk_list_data_point) == 1000 or date == data["dates"][-1]: 

76 SimulationDataPoint.objects.bulk_create(bulk_list_data_point) 

77 bulk_list_data_point = [] 

78 SimulationDataPointExtraInfo.objects.bulk_create(bulk_list_extra_info) 

79 bulk_list_extra_info = [] 

80 

81 return simulation