Coverage for backend/django/PinchAnalysis/views/PinchUtilityViewSet.py: 57%

1from core.viewset import ModelViewSet 

2from PinchAnalysis.serializers.PinchInputSerializers import PinchUtilitySerializer 

3from PinchAnalysis.models.InputModels import PinchInputs, PinchUtility 

4from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes 

5from rest_framework.response import Response 

6from rest_framework import serializers 

7from rest_framework.decorators import action 

8import traceback 

9from PinchAnalysis.models.StreamDataProject import StreamDataProject 

10 

11class BulkCreateUtilitiesSerializer(serializers.Serializer): 

12 projectID = serializers.IntegerField(required=True) 

13 utilities = serializers.ListField( 

14 child=serializers.DictField(), 

15 required=True, 

16 ) 

17 

18class DeleteAllUtilitiesSerializer(serializers.Serializer): 

19 projectID = serializers.IntegerField(required=True) 

20 

21class PinchUtilityViewSet(ModelViewSet): 

22 serializer_class = PinchUtilitySerializer 

23 

24 def get_queryset(self): 

25 queryset = PinchUtility.objects.all() 

26 projectId = self.request.query_params.get("projectOwner") 

27 if projectId is not None: 

28 pinch_inputs = PinchInputs.objects.filter(project_owner_id=projectId) 

29 queryset = queryset.filter(input_owner__in=pinch_inputs) 

30 return queryset 

31 

32 @extend_schema( 

33 parameters=[ 

34 OpenApiParameter(name="projectOwner", required=True, type=OpenApiTypes.INT), 

35 ] 

36 ) 

37 

38 def list(self, request): 

39 return super().list(request) 

40 

41 

42 @extend_schema(request=BulkCreateUtilitiesSerializer, responses=None) 

43 @action(methods=['post'], detail=False, url_path='bulk-create') 

44 def bulk_create(self, request): 

45 try: 

46 serializer = BulkCreateUtilitiesSerializer(data=request.data) 

47 serializer.is_valid(raise_exception=True) 

48 validated_data = serializer.validated_data 

49 

50 projectID = validated_data.get("projectID") 

51 project = StreamDataProject.objects.get(pk=projectID) 

52 utilities = validated_data.get("utilities") 

53 

54 utility_objects = [ 

55 PinchUtility( 

56 input_owner=project.Inputs, 

57 flowsheet_state=project.flowsheet_state, 

58 **stream, 

59 ) 

60 for stream in utilities 

61 ] 

62 

63 PinchUtility.objects.bulk_create(utility_objects) 

64 

65 return Response({'status': 'success'}, status=201) 

66 except Exception as e: 

67 return self.error_response(e) 

68 

69 

70 @extend_schema(request=DeleteAllUtilitiesSerializer, responses=None) 

71 @action(methods=['post'], detail=False, url_path='delete-all') 

72 def delete_all(self, request): 

73 try: 

74 serializer = DeleteAllUtilitiesSerializer(data=request.data) 

75 serializer.is_valid(raise_exception=True) 

76 validated_data = serializer.validated_data 

77 

78 projectID = validated_data.get("projectID") 

79 project = StreamDataProject.objects.get(pk=projectID) 

80 

81 project.Inputs.PinchUtilities.all().delete() 

82 

83 return Response({'status': 'success'}, status=204) 

84 except Exception as e: 

85 return self.error_response(e) 

86 

87 def error_response(self, e): 

88 tb_info = traceback.format_exc() 

89 error_message = str(e) 

90 response_data = {'status': 'error', 'message': error_message, 'traceback': tb_info} 

91 return Response(response_data, status=400)