Coverage for backend/PinchAnalysis/views/PinchUtilityViewSet.py: 42%

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 flowsheet = request.query_params.get("flowsheet") 

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

53 utilities = validated_data.get("utilities") 

54 

55 utility_objects = [ 

56 PinchUtility(input_owner=project.Inputs, flowsheet_id=flowsheet, **stream) for stream in utilities 

57 ] 

58 

59 PinchUtility.objects.bulk_create(utility_objects) 

60 

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

62 except Exception as e: 

63 return self.error_response(e) 

64 

65 

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

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

68 def delete_all(self, request): 

69 try: 

70 serializer = DeleteAllUtilitiesSerializer(data=request.data) 

71 serializer.is_valid(raise_exception=True) 

72 validated_data = serializer.validated_data 

73 

74 projectID = validated_data.get("projectID") 

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

76 

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

78 

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

80 except Exception as e: 

81 return self.error_response(e) 

82 

83 def error_response(self, e): 

84 tb_info = traceback.format_exc() 

85 error_message = str(e) 

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

87 return Response(response_data, status=400)