Coverage for backend/core/auxiliary/views/DownloadMSSData.py: 43%

1from core.auxiliary.models.Expression import Expression 

2from core.auxiliary.serializers.ExpressionSerializer import ExpressionSerializer 

3from drf_spectacular.utils import OpenApiParameter, OpenApiTypes, extend_schema 

4from rest_framework.decorators import api_view, renderer_classes 

5from rest_framework.response import Response 

6from rest_framework_csv.renderers import CSVRenderer 

7from rest_framework import serializers 

8import csv 

9from django.http import HttpResponse 

10from core.validation import api_view_validate 

11from core.auxiliary.methods.export_scenario_data import export_scenario_data, collate 

12from core.auxiliary.models.Scenario import Scenario 

13from core.auxiliary.models.Flowsheet import Flowsheet 

14 

15def transform_to_csv(data): 

16 csv_data = [] 

17 for item in data: 

18 for solveValue in item.get('solveValues'): 

19 csv_data.append({'name': item.get('name'), 'value': solveValue.get('value')}) 

20 return csv_data 

21 

22class MSSRenderer (CSVRenderer): 

23 header = ['name', 'value'] 

24 

25@api_view_validate 

26@extend_schema(parameters=[ 

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

28]) 

29@api_view(['GET']) 

30@renderer_classes([MSSRenderer]) 

31def download_data(request): 

32 """ 

33 This endpoint allows downloading of the input data for a multisteadystate simulation. 

34 I am not sure if this has been updated to reflect the new Scenario model, or if 

35 it is still used in the UI. 

36 """ 

37 flowsheet_id = request.query_params.get("flowsheet") 

38 if not flowsheet_id: 

39 return Response( 

40 {"error": "flowsheetId parameter is required."}, 

41 status=400 

42 ) 

43 expressions = Expression.objects.filter(flowsheet=flowsheet_id) 

44 serialized_data = ExpressionSerializer(expressions, many=True).data 

45 

46 csv_data = transform_to_csv(serialized_data) 

47 

48 response = HttpResponse(content_type='text/csv') 

49 response['Content-Disposition'] = 'attachment; filename="export.csv"' 

50 writer = csv.DictWriter(response, fieldnames=['name', 'value']) 

51 writer.writeheader() 

52 writer.writerows(csv_data) 

53 

54 return response 

55 

56@api_view_validate 

57@extend_schema(parameters=[ 

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

59 OpenApiParameter(name="scenario", required=True, type=OpenApiTypes.INT), 

60]) 

61@api_view(['GET']) 

62@renderer_classes([CSVRenderer]) 

63def download_mss_results(request): 

64 """ 

65 This endpoint allows downloading of the results data for a multisteadystate simulation. 

66 """ 

67 flowsheet_id = request.query_params.get("flowsheet") 

68 scenario_id = request.query_params.get("scenario") 

69 flowsheet = Flowsheet.objects.get(pk=flowsheet_id) 

70 scenario = Scenario.objects.get(pk=scenario_id, flowsheet=flowsheet) 

71 data = export_scenario_data(flowsheet, scenario) 

72 response = HttpResponse(content_type='text/csv') 

73 response['Content-Disposition'] = 'attachment; filename="data.csv"' 

74 writer = csv.DictWriter(response, fieldnames=data.keys()) 

75 writer.writeheader() 

76 writer.writerows(collate(data)) 

77 

78 return response