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

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 

13 

14 

15class CreateFlowsheetTemplateSerializer(serializers.Serializer): 

16 flowsheet_id = serializers.IntegerField() 

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

18 

19 

20class CreateModuleFromTemplateSerializer(serializers.Serializer): 

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

22 flowsheet = serializers.IntegerField() 

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

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

25 

26 

27class FlowsheetTemplateViewSet(viewsets.ReadOnlyModelViewSet): 

28 serializer_class = FlowsheetSerializer 

29 

30 def get_queryset(self): 

31 # Public templates are visible to all users 

32 # Private templates only to the owner 

33 flowsheets = Flowsheet.objects.select_related('owner') 

34 

35 public_flowsheet_templates = flowsheets.filter( 

36 flowsheet_template_type=FlowsheetTemplateType.PublicTemplate) 

37 private_flowsheet_templates = flowsheets.filter( 

38 flowsheet_template_type=FlowsheetTemplateType.PrivateTemplate, 

39 owner=self.request.user 

40 ) 

41 return public_flowsheet_templates | private_flowsheet_templates 

42 

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

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

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

46 """ 

47 Convert a flowsheet to a template 

48 """ 

49 serializer = CreateFlowsheetTemplateSerializer(data=request.data) 

50 serializer.is_valid(raise_exception=True) 

51 

52 flowsheet_id = serializer.validated_data['flowsheet_id'] 

53 flowsheet_template_type = serializer.validated_data['flowsheet_template_type'] 

54 

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

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

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

58 status=status.HTTP_403_FORBIDDEN) 

59 

60 try: 

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

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

63 

64 # update the template type 

65 source_flowsheet.flowsheet_template_type = flowsheet_template_type 

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

67 

68 return Response(FlowsheetSerializer(source_flowsheet).data, status=status.HTTP_200_OK) 

69 

70 except Flowsheet.DoesNotExist: 

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

72 status=status.HTTP_404_NOT_FOUND) 

73 

74 @extend_schema(responses=None) 

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

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

77 """ 

78 Create a new flowsheet from a template 

79 """ 

80 try: 

81 flowsheet_template = self.get_object() 

82 

83 # copy the template to create a new flowsheet 

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

85 

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

87 new_flowsheet.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

88 new_flowsheet.save() 

89 

90 return Response(FlowsheetSerializer(new_flowsheet).data, status=status.HTTP_201_CREATED) 

91 

92 except Flowsheet.DoesNotExist: 

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

94 except Exception as e: 

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

96 

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

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

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

100 """ 

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

102 """ 

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

104 

105 serializer = CreateModuleFromTemplateSerializer(data=data) 

106 serializer.is_valid(raise_exception=True) 

107 

108 current_flowsheet_id = serializer.validated_data['flowsheet'] 

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

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

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

112 

113 try: 

114 with transaction.atomic(): 

115 flowsheet_template = self.get_object() 

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

117 

118 current_group: Grouping | None = None 

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

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

121 

122 # Create the module from template 

123 new_module = create_module_from_template_logic( 

124 flowsheet_template, 

125 current_flowsheet, 

126 request.user, 

127 current_group, 

128 x, 

129 y 

130 ) 

131 

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

133 

134 except Flowsheet.DoesNotExist: 

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

136 except Exception as e: 

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

138 

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

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

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

142 """ 

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

144 """ 

145 try: 

146 flowsheet_template: Flowsheet = self.get_object() 

147 

148 # Must currently be a template 

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

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

151 

152 # Permission checks 

153 if flowsheet_template.flowsheet_template_type == FlowsheetTemplateType.PublicTemplate: 

154 if not request.user.is_staff: 

155 return Response( 

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

157 status=status.HTTP_403_FORBIDDEN 

158 ) 

159 

160 # Revert to regular flowsheet 

161 flowsheet_template.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

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

163 

164 return Response(FlowsheetSerializer(flowsheet_template).data, status=status.HTTP_200_OK) 

165 

166 except Flowsheet.DoesNotExist: 

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

168 status=status.HTTP_404_NOT_FOUND) 

169 except Exception as e: 

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