Coverage for backend/django/core/auxiliary/viewsets/FlowsheetTemplateViewSet.py: 90%

1from django.db import transaction 

2from rest_framework import status, viewsets 

3from rest_framework.decorators import action 

4from rest_framework.response import Response 

5from core.auxiliary.enums.FlowsheetTemplateType import FlowsheetTemplateType 

6from core.auxiliary.models.Flowsheet import Flowsheet 

7from core.auxiliary.serializers.FlowsheetSerializer import FlowsheetSerializer 

8from flowsheetInternals.graphicData.models.groupingModel import Grouping 

9from core.auxiliary.methods.CopyFlowsheet import copy_flowsheet_data, create_module_from_template_logic 

10from flowsheetInternals.graphicData.serializers.groupingSerializer import GroupingSerializer 

11from drf_spectacular.utils import extend_schema 

12from rest_framework import serializers 

13from core.managers import get_flowsheet_access 

14 

15 

16class CreateFlowsheetTemplateSerializer(serializers.Serializer): 

17 flowsheet_id = serializers.IntegerField() 

18 flowsheet_template_type = serializers.ChoiceField(choices=FlowsheetTemplateType.choices) 

19 

20 

21class CreateModuleFromTemplateSerializer(serializers.Serializer): 

22 current_group = serializers.IntegerField(required=False, allow_null=True) 

23 flowsheet = serializers.IntegerField() 

24 x = serializers.FloatField(required=False, allow_null=True) 

25 y = serializers.FloatField(required=False, allow_null=True) 

26 

27 

28class FlowsheetTemplateViewSet(viewsets.ReadOnlyModelViewSet): 

29 serializer_class = FlowsheetSerializer 

30 

31 def get_queryset(self): 

32 # Public templates are visible to all users 

33 # Private templates only to the owner 

34 flowsheets = Flowsheet.with_user_access( 

35 Flowsheet.objects.select_related('owner'), 

36 self.request.user, 

37 ) 

38 

39 public_flowsheet_templates = flowsheets.filter( 

40 flowsheet_template_type=FlowsheetTemplateType.PublicTemplate) 

41 private_flowsheet_templates = flowsheets.filter( 

42 flowsheet_template_type=FlowsheetTemplateType.PrivateTemplate, 

43 owner=self.request.user 

44 ) 

45 return public_flowsheet_templates | private_flowsheet_templates 

46 

47 @extend_schema(request=CreateFlowsheetTemplateSerializer, responses=None) 

48 @action(detail=False, methods=['post'], url_path='create-flowsheet-template') 

49 def create_flowsheet_template(self, request) -> Response: 

50 """ 

51 Convert a flowsheet to a template 

52 """ 

53 serializer = CreateFlowsheetTemplateSerializer(data=request.data) 

54 serializer.is_valid(raise_exception=True) 

55 

56 flowsheet_id = serializer.validated_data['flowsheet_id'] 

57 flowsheet_template_type = serializer.validated_data['flowsheet_template_type'] 

58 

59 # Only allow conversion to public template if user is staff 

60 if flowsheet_template_type == FlowsheetTemplateType.PublicTemplate and not request.user.is_staff: 

61 return Response({'error': 'You do not have permission to create public templates'}, 

62 status=status.HTTP_403_FORBIDDEN) 

63 

64 try: 

65 # get the source flowsheet (only owner can convert to template) 

66 source_flowsheet = Flowsheet.objects.select_related('owner').get(id=flowsheet_id, owner=request.user) 

67 

68 # update the template type 

69 source_flowsheet.flowsheet_template_type = flowsheet_template_type 

70 source_flowsheet.save(update_fields=['flowsheet_template_type']) 

71 

72 return Response( 

73 FlowsheetSerializer(source_flowsheet, context={"request": request}).data, 

74 status=status.HTTP_200_OK, 

75 ) 

76 

77 except Flowsheet.DoesNotExist: 

78 return Response({'error': 'Flowsheet not found or you do not have permission'}, 

79 status=status.HTTP_404_NOT_FOUND) 

80 

81 @extend_schema(responses=None) 

82 @action(detail=True, methods=['post'], url_path='create-from-flowsheet-template') 

83 def create_from_flowsheet_template(self, request, pk=None) -> Response: 

84 """ 

85 Create a new flowsheet from a template 

86 """ 

87 try: 

88 flowsheet_template = self.get_object() 

89 

90 # copy the template to create a new flowsheet 

91 new_flowsheet = copy_flowsheet_data(flowsheet_template, user=request.user) 

92 

93 # reset new flowsheet's template type to regular flowsheet 

94 new_flowsheet.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

95 new_flowsheet.save() 

96 

97 return Response( 

98 FlowsheetSerializer(new_flowsheet, context={"request": request}).data, 

99 status=status.HTTP_201_CREATED, 

100 ) 

101 

102 except Flowsheet.DoesNotExist: 

103 return Response({'error': 'Template not found'}, status=status.HTTP_404_NOT_FOUND) 

104 except Exception as e: 

105 return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST) 

106 

107 @extend_schema(request=CreateModuleFromTemplateSerializer, responses=None) 

108 @action(detail=True, methods=['post'], url_path='create-module-from-template') 

109 def create_module_from_template(self, request, pk=None) -> Response: 

110 """ 

111 Create a new module (grouping) from a template within the current flowsheet 

112 """ 

113 data = request.data.get('createModuleFromTemplate', request.data) 

114 

115 serializer = CreateModuleFromTemplateSerializer(data=data) 

116 serializer.is_valid(raise_exception=True) 

117 

118 current_flowsheet_id = serializer.validated_data['flowsheet'] 

119 current_group_id = serializer.validated_data.get('current_group') 

120 x = serializer.validated_data.get("x") 

121 y = serializer.validated_data.get("y") 

122 

123 try: 

124 with transaction.atomic(): 

125 flowsheet_template = self.get_object() 

126 access_state = get_flowsheet_access(request.user, current_flowsheet_id) 

127 if not access_state.has_read_access: 

128 return Response({'error': 'Template or flowsheet not found'}, status=status.HTTP_404_NOT_FOUND) 

129 if not access_state.has_write_access: 

130 return Response( 

131 {'error': 'You do not have write access to this flowsheet'}, 

132 status=status.HTTP_403_FORBIDDEN, 

133 ) 

134 

135 current_flowsheet = Flowsheet.objects.get(id=current_flowsheet_id) 

136 

137 current_group: Grouping | None = None 

138 if current_group_id: 138 ↛ 142line 138 didn't jump to line 142 because the condition on line 138 was always true

139 current_group = Grouping.objects.get(id=current_group_id) 

140 

141 # Create the module from template 

142 new_module = create_module_from_template_logic( 

143 flowsheet_template, 

144 current_flowsheet, 

145 request.user, 

146 current_group, 

147 x, 

148 y 

149 ) 

150 

151 return Response(GroupingSerializer(new_module).data, status=status.HTTP_201_CREATED) 

152 

153 except Flowsheet.DoesNotExist: 

154 return Response({'error': 'Template or flowsheet not found'}, status=status.HTTP_404_NOT_FOUND) 

155 except Exception as e: 

156 return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST) 

157 

158 @extend_schema(request=None, responses=FlowsheetSerializer) 

159 @action(detail=True, methods=['post'], url_path='revert-flowsheet-template') 

160 def revert_flowsheet_template(self, request, pk=None) -> Response: 

161 """ 

162 Revert a template flowsheet back to a regular (not template) flowsheet. 

163 """ 

164 try: 

165 flowsheet_template: Flowsheet = self.get_object() 

166 

167 # Must currently be a template 

168 if flowsheet_template.flowsheet_template_type == FlowsheetTemplateType.NotTemplate: 168 ↛ 169line 168 didn't jump to line 169 because the condition on line 168 was never true

169 return Response({'error': 'Flowsheet is not a template'}, status=status.HTTP_400_BAD_REQUEST) 

170 

171 # Permission checks 

172 if flowsheet_template.flowsheet_template_type == FlowsheetTemplateType.PublicTemplate: 

173 if not request.user.is_staff: 

174 return Response( 

175 {'error': 'You do not have permission to revert public templates'}, 

176 status=status.HTTP_403_FORBIDDEN 

177 ) 

178 

179 # Revert to regular flowsheet 

180 flowsheet_template.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

181 flowsheet_template.save(update_fields=['flowsheet_template_type']) 

182 

183 return Response( 

184 FlowsheetSerializer(flowsheet_template, context={"request": request}).data, 

185 status=status.HTTP_200_OK, 

186 ) 

187 

188 except Flowsheet.DoesNotExist: 

189 return Response({'error': 'Template not found or you do not have permission'}, 

190 status=status.HTTP_404_NOT_FOUND) 

191 except Exception as e: 

192 return Response({'error': str(e)}, status=status.HTTP_400_BAD_REQUEST)