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

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 

19class CreateModuleFromTemplateSerializer(serializers.Serializer): 

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

21 flowsheet = serializers.IntegerField() 

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

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

24 

25class FlowsheetTemplateViewSet(viewsets.ReadOnlyModelViewSet): 

26 serializer_class = FlowsheetSerializer 

27 

28 def get_queryset(self): 

29 # Public templates are visible to all users 

30 # Private templates only to the owner 

31 public_flowsheet_templates = Flowsheet.objects.filter(flowsheet_template_type=FlowsheetTemplateType.PublicTemplate) 

32 private_flowsheet_templates = Flowsheet.objects.filter( 

33 flowsheet_template_type=FlowsheetTemplateType.PrivateTemplate, 

34 owner=self.request.user 

35 ) 

36 return public_flowsheet_templates | private_flowsheet_templates 

37 

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

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

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

41 """ 

42 Convert a flowsheet to a template 

43 """ 

44 serializer = CreateFlowsheetTemplateSerializer(data=request.data) 

45 serializer.is_valid(raise_exception=True) 

46 

47 flowsheet_id = serializer.validated_data['flowsheet_id'] 

48 flowsheet_template_type = serializer.validated_data['flowsheet_template_type'] 

49 

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

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

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

53 status=status.HTTP_403_FORBIDDEN) 

54 

55 try: 

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

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

58 

59 # update the template type 

60 source_flowsheet.flowsheet_template_type = flowsheet_template_type 

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

62 

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

64 

65 except Flowsheet.DoesNotExist: 

66 return Response({'error': 'Flowsheet not found or you do not have permission'}, status=status.HTTP_404_NOT_FOUND) 

67 

68 @extend_schema(responses=None) 

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

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

71 """ 

72 Create a new flowsheet from a template 

73 """ 

74 try: 

75 flowsheet_template = self.get_object() 

76 

77 # copy the template to create a new flowsheet 

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

79 

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

81 new_flowsheet.flowsheet_template_type = FlowsheetTemplateType.NotTemplate 

82 new_flowsheet.save() 

83 

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

85 

86 except Flowsheet.DoesNotExist: 

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

88 except Exception as e: 

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

90 

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

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

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

94 """ 

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

96 """ 

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

98 

99 serializer = CreateModuleFromTemplateSerializer(data=data) 

100 serializer.is_valid(raise_exception=True) 

101 

102 current_flowsheet_id = serializer.validated_data['flowsheet'] 

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

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

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

106 

107 

108 try: 

109 with transaction.atomic(): 

110 flowsheet_template = self.get_object() 

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

112 

113 current_group: Grouping | None = None 

114 if current_group_id: 

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

116 

117 # Create the module from template 

118 new_module = create_module_from_template_logic( 

119 flowsheet_template, 

120 current_flowsheet, 

121 request.user, 

122 current_group, 

123 x, 

124 y 

125 ) 

126 

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

128 

129 except Flowsheet.DoesNotExist: 

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

131 except Exception as e: 

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

133 

134