Coverage for backend/django/Economics/costing/costable_items/viewsets.py: 92%

96 statements  

« prev     ^ index     » next       coverage.py v7.10.7, created at 2026-07-22 05:22 +0000

1from django.db import transaction 

2from drf_spectacular.utils import extend_schema 

3from rest_framework.decorators import action 

4from rest_framework.exceptions import NotFound, ValidationError 

5from rest_framework.response import Response 

6 

7from core.viewset import ModelViewSet 

8from Economics.costing.models import CapitalCostLine, CostCurve, CostDriver, CostableItem, EquipmentMapping 

9from Economics.costing.costable_items.serializers import ( 

10 BulkEquipmentSetupRequestSerializer, 

11 BulkEquipmentSetupResultSerializer, 

12 CostDriverPropertyOptionSerializer, 

13 CostDriverSerializer, 

14 CostableItemSerializer, 

15 EquipmentMappingSerializer, 

16) 

17from Economics.costing.cost_curves.driver_properties import ( 

18 apply_bulk_driver_inputs, 

19 cost_driver_property_options, 

20 preview_or_apply_bulk_equipment_setup, 

21) 

22from Economics.costing.capital.generated_lines import ( 

23 sync_generated_capital_lines, 

24 sync_generated_capital_lines_for_mappings, 

25) 

26from Economics.costing.capital.capital_line_sources import GENERATED_CAPITAL_LINE_SOURCE 

27from Economics.shared.access import require_write_access 

28from Economics.studies.api_mutations import StudyMutationMixin 

29from Economics.studies.services.result_staleness import mark_study_stale 

30from Economics.studies.models import EconomicsStudy 

31 

32 

33class CostableItemViewSet(StudyMutationMixin, ModelViewSet): 

34 serializer_class = CostableItemSerializer 

35 stale_reason = "economics_costable_item_api_saved" 

36 sync_generated_capital_lines = True 

37 

38 def get_queryset(self): 

39 queryset = CostableItem.objects.select_related("study", "simulation_object").prefetch_related("cost_driver", "equipment_mapping") 

40 queryset = queryset.exclude(simulation_object__objectType="group") 

41 study = self.request.query_params.get("study") 

42 return queryset.filter(study_id=study) if study is not None else queryset 

43 

44 

45class CostDriverViewSet(StudyMutationMixin, ModelViewSet): 

46 serializer_class = CostDriverSerializer 

47 stale_reason = "economics_cost_driver_api_saved" 

48 sync_generated_capital_lines = True 

49 

50 def _study_for_instance(self, instance) -> EconomicsStudy | None: 

51 return instance.costable_item.study if instance.costable_item_id else None 

52 

53 def get_queryset(self): 

54 queryset = CostDriver.objects.select_related("costable_item__study", "property_info", "manual_property_info") 

55 costable_item = self.request.query_params.get("costable_item") 

56 return queryset.filter(costable_item_id=costable_item) if costable_item is not None else queryset 

57 

58 @extend_schema(responses=CostDriverPropertyOptionSerializer(many=True)) 

59 @action(detail=True, methods=["get"], url_path="property-options") 

60 def property_options(self, request, pk=None): 

61 driver = self.get_object() 

62 serializer = CostDriverPropertyOptionSerializer(cost_driver_property_options(driver), many=True) 

63 return Response(serializer.data) 

64 

65 

66class EquipmentMappingViewSet(StudyMutationMixin, ModelViewSet): 

67 serializer_class = EquipmentMappingSerializer 

68 stale_reason = "economics_equipment_mapping_api_saved" 

69 sync_generated_capital_lines = True 

70 

71 def _study_for_instance(self, instance) -> EconomicsStudy | None: 

72 return instance.costable_item.study if instance.costable_item_id else None 

73 

74 def get_queryset(self): 

75 queryset = EquipmentMapping.objects.select_related( 

76 "costable_item__study", 

77 "costable_item__study__settings_profile", 

78 "costable_item__study__schedule_scenario", 

79 "costable_item__cost_driver", 

80 "costable_item__simulation_object", 

81 "cost_curve", 

82 ) 

83 costable_item = self.request.query_params.get("costable_item") 

84 return queryset.filter(costable_item_id=costable_item) if costable_item is not None else queryset 

85 

86 def perform_update(self, serializer): 

87 """Refresh only the generated line owned by the updated mapping.""" 

88 

89 instance = serializer.save() 

90 study = instance.costable_item.study 

91 sync_generated_capital_lines_for_mappings([instance]) 

92 mark_study_stale(study, reason=self.stale_reason) 

93 

94 @extend_schema( 

95 request=BulkEquipmentSetupRequestSerializer, 

96 responses=BulkEquipmentSetupResultSerializer(many=True), 

97 ) 

98 @action(detail=False, methods=["post"], url_path="bulk-setup") 

99 def bulk_setup(self, request): 

100 request_serializer = BulkEquipmentSetupRequestSerializer(data=request.data) 

101 request_serializer.is_valid(raise_exception=True) 

102 payload = request_serializer.validated_data 

103 dry_run = payload["dry_run"] 

104 if not dry_run: 

105 require_write_access(request.user) 

106 

107 mapping_ids = payload["mapping_ids"] 

108 mappings_by_id = self.get_queryset().in_bulk(mapping_ids) 

109 missing_mapping_ids = [mapping_id for mapping_id in mapping_ids if mapping_id not in mappings_by_id] 

110 if missing_mapping_ids: 110 ↛ 111line 110 didn't jump to line 111 because the condition on line 110 was never true

111 raise NotFound("One or more equipment mappings were not found.") 

112 mappings = [mappings_by_id[mapping_id] for mapping_id in mapping_ids] 

113 study_ids = {mapping.costable_item.study_id for mapping in mappings} 

114 if len(study_ids) != 1: 114 ↛ 115line 114 didn't jump to line 115 because the condition on line 114 was never true

115 raise ValidationError({"mapping_ids": "Bulk setup can only update one economics study at a time."}) 

116 

117 cost_curve = None 

118 cost_curve_id = payload.get("cost_curve") 

119 if cost_curve_id is not None: 119 ↛ 124line 119 didn't jump to line 124 because the condition on line 119 was always true

120 try: 

121 cost_curve = CostCurve.objects.get(pk=cost_curve_id) 

122 except CostCurve.DoesNotExist as error: 

123 raise NotFound("Selected cost curve was not found.") from error 

124 if payload["apply_cost_curve"] and cost_curve is not None: 124 ↛ 135line 124 didn't jump to line 135 because the condition on line 124 was always true

125 for mapping in mappings: 

126 serializer = EquipmentMappingSerializer( 

127 mapping, 

128 context=self.get_serializer_context(), 

129 ) 

130 # Relations are already authorized model instances from the 

131 # scoped bulk query, so run semantic validation without 

132 # deserializing the same curve PK once per mapping. 

133 serializer.validate({"cost_curve": cost_curve}) 

134 

135 with transaction.atomic(): 

136 results = preview_or_apply_bulk_equipment_setup( 

137 mappings=mappings, 

138 cost_curve=cost_curve, 

139 dry_run=dry_run, 

140 apply_cost_curve=payload["apply_cost_curve"], 

141 apply_recommended_sizing_property=payload["apply_recommended_sizing_property"], 

142 overwrite_sizing_property=payload["overwrite_sizing_property"], 

143 driver_inputs=payload["driver_inputs"], 

144 ) 

145 if not dry_run: 

146 study = mappings[0].costable_item.study 

147 generated_item_ids = set( 

148 CapitalCostLine.objects.filter( 

149 costable_item_id__in=[mapping.costable_item_id for mapping in mappings], 

150 source=GENERATED_CAPITAL_LINE_SOURCE, 

151 ).values_list("costable_item_id", flat=True) 

152 ) 

153 missing_line_mappings = [ 

154 mapping 

155 for mapping in mappings 

156 if mapping.costable_item_id not in generated_item_ids 

157 ] 

158 if missing_line_mappings: 

159 sync_generated_capital_lines_for_mappings(missing_line_mappings) 

160 apply_bulk_driver_inputs( 

161 mappings=mappings, 

162 cost_curve=cost_curve, 

163 apply_cost_curve=payload["apply_cost_curve"], 

164 apply_recommended_sizing_property=payload["apply_recommended_sizing_property"], 

165 overwrite_sizing_property=payload["overwrite_sizing_property"], 

166 driver_inputs=payload["driver_inputs"], 

167 setup_results=results, 

168 ) 

169 sync_generated_capital_lines_for_mappings(mappings) 

170 mark_study_stale(study, reason="economics_equipment_bulk_setup_saved") 

171 

172 return Response(BulkEquipmentSetupResultSerializer(results, many=True).data)