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 

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 public_flowsheet_templates = Flowsheet.objects.filter( 

34 flowsheet_template_type=FlowsheetTemplateType.PublicTemplate) 

35 private_flowsheet_templates = Flowsheet.objects.filter( 

36 flowsheet_template_type=FlowsheetTemplateType.PrivateTemplate, 

37 owner=self.request.user 

38 ) 

39 return public_flowsheet_templates | private_flowsheet_templates 

40 

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

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

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

44 """ 

45 Convert a flowsheet to a template 

46 """ 

47 serializer = CreateFlowsheetTemplateSerializer(data=request.data) 

48 serializer.is_valid(raise_exception=True) 

49 

50 flowsheet_id = serializer.validated_data['flowsheet_id'] 

51 flowsheet_template_type = serializer.validated_data['flowsheet_template_type'] 

52 

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

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

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

56 status=status.HTTP_403_FORBIDDEN) 

57 

58 try: 

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

60 source_flowsheet = Flowsheet.objects.get(id=flowsheet_id, owner=request.user) 

61 

62 # update the template type 

63 source_flowsheet.flowsheet_template_type = flowsheet_template_type 

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

65 

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

67 

68 except Flowsheet.DoesNotExist: 

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

70 status=status.HTTP_404_NOT_FOUND) 

71 

72 @extend_schema(responses=None) 

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

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

75 """ 

76 Create a new flowsheet from a template 

77 """ 

78 try: 

79 flowsheet_template = self.get_object() 

80 

81 # copy the template to create a new flowsheet 

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

83 

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

85 new_flowsheet.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

86 new_flowsheet.save() 

87 

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

89 

90 except Flowsheet.DoesNotExist: 

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

92 except Exception as e: 

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

94 

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

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

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

98 """ 

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

100 """ 

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

102 

103 serializer = CreateModuleFromTemplateSerializer(data=data) 

104 serializer.is_valid(raise_exception=True) 

105 

106 current_flowsheet_id = serializer.validated_data['flowsheet'] 

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

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

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

110 

111 try: 

112 with transaction.atomic(): 

113 flowsheet_template = self.get_object() 

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

115 

116 current_group: Grouping | None = None 

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

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

119 

120 # Create the module from template 

121 new_module = create_module_from_template_logic( 

122 flowsheet_template, 

123 current_flowsheet, 

124 request.user, 

125 current_group, 

126 x, 

127 y 

128 ) 

129 

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

131 

132 except Flowsheet.DoesNotExist: 

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

134 except Exception as e: 

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

136 

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

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

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

140 """ 

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

142 """ 

143 try: 

144 flowsheet_template: Flowsheet = self.get_object() 

145 

146 # Must currently be a template 

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

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

149 

150 # Permission checks 

151 if flowsheet_template.flowsheet_template_type == FlowsheetTemplateType.PublicTemplate: 

152 if not request.user.is_staff: 152 ↛ 159line 152 didn't jump to line 159 because the condition on line 152 was always true

153 return Response( 

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

155 status=status.HTTP_403_FORBIDDEN 

156 ) 

157 

158 # Revert to regular flowsheet 

159 flowsheet_template.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

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

161 

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

163 

164 except Flowsheet.DoesNotExist: 

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

166 status=status.HTTP_404_NOT_FOUND) 

167 except Exception as e: 

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