Coverage for backend/PinchAnalysis/views/StreamDataProjectViewSet.py: 44%

1import traceback 

2from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes 

3from core.viewset import ModelViewSet 

4from rest_framework.response import Response 

5from PinchAnalysis.models.StreamDataProject import StreamDataProject 

6from PinchAnalysis.serializers.StreamDataProjectSerializer import StreamDataProjectSerializer 

7from flowsheetInternals.graphicData.models.groupingModel import Grouping 

8class StreamDataProjectViewSet(ModelViewSet): 

9 serializer_class = StreamDataProjectSerializer 

10 

11 def get_queryset(self): 

12 return StreamDataProject.objects.all() 

13 

14 @extend_schema( 

15 parameters=[ 

16 OpenApiParameter(name="flowsheet", required=True, type=OpenApiTypes.INT), 

17 ] 

18 ) 

19 

20 def list(self, request): 

21 return super().list(request) 

22 

23 def error_response(self, e): 

24 tb_info = traceback.format_exc() 

25 error_message = str(e) 

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

27 return Response(response_data, status=400) 

28 

29 def create(self, request, *args, **kwargs): 

30 try: 

31 group_owner = request.data.get('groupOwner') 

32 if not group_owner: 

33 return Response( 

34 {'error': 'groupOwner is required.'}, 

35 status=400 

36 ) 

37 grouping_owner = Grouping.objects.get(pk=group_owner) 

38 flowsheet = grouping_owner.simulationObject.flowsheet 

39 instance = StreamDataProject.create(flowsheet=flowsheet) 

40 return Response({'status': 'success'}, status=200) 

41 except Exception as e: 

42 return self.error_response(e)