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

88 statements  

« prev     ^ index     » next       coverage.py v7.10.7, created at 2026-06-23 21:51 +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 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 _sync_generated_capital_lines 

23from Economics.shared.access import _require_write_access 

24from Economics.studies.api_mutations import StudyMutationMixin, _mark_study_stale 

25from Economics.studies.models import EconomicsStudy 

26 

27 

28class CostableItemViewSet(StudyMutationMixin, ModelViewSet): 

29 serializer_class = CostableItemSerializer 

30 stale_reason = "economics_costable_item_api_saved" 

31 sync_generated_capital_lines = True 

32 

33 def get_queryset(self): 

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

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

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

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

38 

39 

40class CostDriverViewSet(StudyMutationMixin, ModelViewSet): 

41 serializer_class = CostDriverSerializer 

42 stale_reason = "economics_cost_driver_api_saved" 

43 sync_generated_capital_lines = True 

44 

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

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

47 

48 def get_queryset(self): 

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

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

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

52 

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

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

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

56 driver = self.get_object() 

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

58 return Response(serializer.data) 

59 

60 

61class EquipmentMappingViewSet(StudyMutationMixin, ModelViewSet): 

62 serializer_class = EquipmentMappingSerializer 

63 stale_reason = "economics_equipment_mapping_api_saved" 

64 sync_generated_capital_lines = True 

65 

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

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

68 

69 def get_queryset(self): 

70 queryset = EquipmentMapping.objects.select_related( 

71 "costable_item__study", 

72 "costable_item__cost_driver", 

73 "costable_item__simulation_object", 

74 "cost_curve", 

75 ) 

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

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

78 

79 @extend_schema( 

80 request=BulkEquipmentSetupRequestSerializer, 

81 responses=BulkEquipmentSetupResultSerializer(many=True), 

82 ) 

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

84 def bulk_setup(self, request): 

85 request_serializer = BulkEquipmentSetupRequestSerializer(data=request.data) 

86 request_serializer.is_valid(raise_exception=True) 

87 payload = request_serializer.validated_data 

88 dry_run = payload["dry_run"] 

89 if not dry_run: 

90 _require_write_access(request.user) 

91 

92 mapping_ids = payload["mapping_ids"] 

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

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

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

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

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

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

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

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

101 

102 cost_curve = None 

103 cost_curve_id = payload.get("cost_curve") 

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

105 try: 

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

107 except CostCurve.DoesNotExist as error: 

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

109 if cost_curve.flowsheet_id != mappings[0].flowsheet_id: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true

110 raise ValidationError({"cost_curve": "Selected cost curve must belong to the same flowsheet."}) 

111 

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

113 for mapping in mappings: 

114 serializer = EquipmentMappingSerializer( 

115 mapping, 

116 data={"cost_curve": cost_curve.pk}, 

117 partial=True, 

118 context=self.get_serializer_context(), 

119 ) 

120 serializer.is_valid(raise_exception=True) 

121 

122 with transaction.atomic(): 

123 results = preview_or_apply_bulk_equipment_setup( 

124 mappings=mappings, 

125 cost_curve=cost_curve, 

126 dry_run=dry_run, 

127 apply_cost_curve=payload["apply_cost_curve"], 

128 apply_recommended_sizing_property=payload["apply_recommended_sizing_property"], 

129 overwrite_sizing_property=payload["overwrite_sizing_property"], 

130 driver_inputs=payload["driver_inputs"], 

131 ) 

132 if not dry_run: 

133 study = mappings[0].costable_item.study 

134 _sync_generated_capital_lines(study) 

135 apply_bulk_driver_inputs( 

136 mappings=mappings, 

137 cost_curve=cost_curve, 

138 apply_cost_curve=payload["apply_cost_curve"], 

139 apply_recommended_sizing_property=payload["apply_recommended_sizing_property"], 

140 overwrite_sizing_property=payload["overwrite_sizing_property"], 

141 driver_inputs=payload["driver_inputs"], 

142 ) 

143 _sync_generated_capital_lines(study) 

144 _mark_study_stale(study, reason="economics_equipment_bulk_setup_saved") 

145 

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