Coverage for backend/django/Economics/costing/cost_curves/viewsets.py: 100%

1from django.db import transaction 

2from drf_spectacular.utils import OpenApiParameter, extend_schema 

3from rest_framework import status 

4from rest_framework.decorators import action 

5from rest_framework.response import Response 

6 

7from core.viewset import ModelViewSet 

8from core.validation import flowsheet_context, get_current_flowsheet 

9from Economics.costing.models import CapitalCostLine, CostCurve, CostCurveTemplate, EquipmentMapping 

10from Economics.costing.cost_curves.catalog import cost_curve_equipment_categories 

11from Economics.costing.cost_curves.serializers import ( 

12 CostCurveAuthoringSerializer, 

13 CostCurveEquipmentCategorySerializer, 

14 CostCurveSerializer, 

15) 

16from Economics.costing.cost_curves.library import save_private_template_from_curve 

17from Economics.costing.capital.generated_lines import sync_generated_capital_lines 

18from Economics.shared.access import require_write_access 

19from Economics.studies.services.result_staleness import mark_study_stale 

20from Economics.studies.models import EconomicsStudy 

21 

22 

23class CostCurveViewSet(ModelViewSet): 

24 serializer_class = CostCurveSerializer 

25 

26 def get_serializer_class(self): 

27 if self.action in {"create", "update", "partial_update"}: 

28 return CostCurveAuthoringSerializer 

29 return CostCurveSerializer 

30 

31 @extend_schema( 

32 parameters=[ 

33 OpenApiParameter("flowsheet", int, OpenApiParameter.QUERY, required=True), 

34 OpenApiParameter("active", bool, OpenApiParameter.QUERY, required=False), 

35 ], 

36 responses=CostCurveSerializer(many=True), 

37 ) 

38 def list(self, request, *args, **kwargs): 

39 return super().list(request, *args, **kwargs) 

40 

41 @extend_schema(request=CostCurveAuthoringSerializer, responses={status.HTTP_201_CREATED: CostCurveSerializer}) 

42 def create(self, request, *args, **kwargs): 

43 require_write_access(request.user) 

44 with transaction.atomic(): 

45 serializer = self.get_serializer(data=request.data) 

46 serializer.is_valid(raise_exception=True) 

47 self.perform_create(serializer) 

48 instance = serializer.instance 

49 headers = self.get_success_headers(serializer.data) 

50 self._save_private_template_if_requested(instance) 

51 response_data = CostCurveSerializer(instance, context=self.get_serializer_context()).data 

52 return Response(response_data, status=status.HTTP_201_CREATED, headers=headers) 

53 

54 @extend_schema(request=CostCurveAuthoringSerializer, responses=CostCurveSerializer) 

55 def update(self, request, *args, **kwargs): 

56 require_write_access(request.user) 

57 with transaction.atomic(): 

58 response = super().update(request, *args, **kwargs) 

59 instance = self.get_object() 

60 self._save_private_template_if_requested(instance) 

61 _refresh_dependent_studies_for_curve(instance) 

62 response.data = CostCurveSerializer(instance, context=self.get_serializer_context()).data 

63 return response 

64 

65 @extend_schema(request=CostCurveAuthoringSerializer(partial=True), responses=CostCurveSerializer) 

66 def partial_update(self, request, *args, **kwargs): 

67 require_write_access(request.user) 

68 with transaction.atomic(): 

69 response = super().partial_update(request, *args, **kwargs) 

70 instance = self.get_object() 

71 self._save_private_template_if_requested(instance) 

72 _refresh_dependent_studies_for_curve(instance) 

73 response.data = CostCurveSerializer(instance, context=self.get_serializer_context()).data 

74 return response 

75 

76 def destroy(self, request, *args, **kwargs): 

77 require_write_access(request.user) 

78 instance = self.get_object() 

79 affected_study_ids = set( 

80 EquipmentMapping._base_manager.filter(cost_curve=instance).values_list( 

81 "costable_item__study_id", flat=True 

82 ) 

83 ) 

84 affected_study_ids.update( 

85 CapitalCostLine._base_manager.filter(cost_curve=instance).values_list("study_id", flat=True) 

86 ) 

87 affected_studies = list( 

88 EconomicsStudy._base_manager.filter(pk__in=affected_study_ids).select_related( 

89 "flowsheet_state__flowsheet__project" 

90 ) 

91 ) 

92 with transaction.atomic(): 

93 response = super().destroy(request, *args, **kwargs) 

94 for study in affected_studies: 

95 _refresh_dependent_study(study) 

96 return response 

97 

98 @extend_schema( 

99 parameters=[ 

100 OpenApiParameter("flowsheet", int, OpenApiParameter.QUERY, required=True), 

101 OpenApiParameter("active", bool, OpenApiParameter.QUERY, required=False), 

102 ], 

103 responses=CostCurveEquipmentCategorySerializer(many=True), 

104 ) 

105 @action(detail=False, methods=["get"], url_path="equipment-options") 

106 def equipment_options(self, request): 

107 options = cost_curve_equipment_categories( 

108 self.get_queryset(), 

109 user_templates=CostCurveTemplate.objects.all(), 

110 ) 

111 serializer = CostCurveEquipmentCategorySerializer(options, many=True) 

112 return Response(serializer.data) 

113 

114 def get_queryset(self): 

115 queryset = CostCurve.objects.all() 

116 active = self.request.query_params.get("active") 

117 return queryset.filter(active=active.lower() == "true") if active is not None else queryset 

118 

119 def perform_create(self, serializer): 

120 self._save_as_template_requested = serializer.validated_data.pop("save_as_template", False) 

121 serializer.save() 

122 

123 def perform_update(self, serializer): 

124 self._save_as_template_requested = serializer.validated_data.pop("save_as_template", False) 

125 serializer.save() 

126 

127 def _save_private_template_if_requested(self, curve: CostCurve) -> None: 

128 if getattr(self, "_save_as_template_requested", False): 

129 save_private_template_from_curve(curve=curve, owner=self.request.user) 

130 

131 

132def _dependent_studies_for_curve(curve: CostCurve) -> list[EconomicsStudy]: 

133 affected_study_ids = set( 

134 EquipmentMapping._base_manager.filter(cost_curve=curve).values_list( 

135 "costable_item__study_id", flat=True 

136 ) 

137 ) 

138 affected_study_ids.update( 

139 CapitalCostLine._base_manager.filter(cost_curve=curve).values_list("study_id", flat=True) 

140 ) 

141 return list( 

142 EconomicsStudy._base_manager.filter(pk__in=affected_study_ids) 

143 .select_related("flowsheet_state__flowsheet__project") 

144 .order_by("pk") 

145 ) 

146 

147 

148def _refresh_dependent_studies_for_curve(curve: CostCurve) -> None: 

149 for study in _dependent_studies_for_curve(curve): 

150 _refresh_dependent_study(study) 

151 

152 

153def _refresh_dependent_study(study: EconomicsStudy) -> None: 

154 """Refresh a referenced study under its own flowsheet access context.""" 

155 flowsheet = study.flowsheet_state.flowsheet 

156 with flowsheet_context(flowsheet.pk, flowsheet.project.owner): 

157 get_current_flowsheet()["bypass_write_checks"] = True 

158 sync_generated_capital_lines(study) 

159 mark_study_stale(study, reason="cost_curve_changed")