Coverage for backend/django/core/auxiliary/viewsets/DataColumnViewSet.py: 87%

1from core.viewset import ModelViewSet 

2from rest_framework.response import Response 

3from core.auxiliary.models.DataColumn import DataColumn 

4from core.auxiliary.serializers.DataColumnSerializer import DataColumnSerializer 

5from core.auxiliary.services.mss_input_data import mark_mss_input_scenarios_changed 

6from drf_spectacular.utils import extend_schema, OpenApiParameter 

7from drf_spectacular.types import OpenApiTypes 

8from rest_framework.decorators import action 

9from rest_framework.serializers import DictField, IntegerField 

10from django.db.models import QuerySet 

11class DataColumnViewSet(ModelViewSet): 

12 serializer_class = DataColumnSerializer 

13 

14 def get_queryset(self): 

15 queryset = DataColumn.objects.select_related( 

16 "property_value__property__set__simulationObject", 

17 ) 

18 return queryset 

19 

20 @extend_schema( 

21 parameters=[ 

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

23 ] 

24 ) 

25 def list(self, request): 

26 queryset = self.get_queryset() 

27 return Response(DataColumnSerializer(queryset.filter(scenario_id=request.query_params.get("scenario")), many=True).data) 

28 

29 def perform_create(self, serializer): 

30 instance = serializer.save() 

31 mark_mss_input_scenarios_changed(scenario_ids=[instance.scenario_id]) 

32 

33 def perform_update(self, serializer): 

34 previous_scenario_id = serializer.instance.scenario_id 

35 instance = serializer.save() 

36 mark_mss_input_scenarios_changed(scenario_ids=[previous_scenario_id, instance.scenario_id]) 

37 

38 def perform_destroy(self, instance): 

39 scenario_id = instance.scenario_id 

40 super().perform_destroy(instance) 

41 mark_mss_input_scenarios_changed(scenario_ids=[scenario_id]) 

42 

43 

44 @extend_schema(responses=DictField(child=IntegerField()), 

45 parameters=[ 

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

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

48 ]) 

49 @action(detail=False,methods=["get"]) 

50 def download_tag_mappings(self, request): 

51 queryset = self.get_queryset().filter( 

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

53 ) 

54 return Response( 

55 create_data_column_json(queryset), 

56 content_type="text/json", 

57 headers={'Content-Disposition': 'attachment; filename="mss-tag-mappings.json"'} 

58 ) 

59 

60def create_data_column_json(queryset: QuerySet[DataColumn, DataColumn]) -> dict[str, dict[str, str]]: 

61 """ 

62 Create a JSON representation of the MSS columns in the scenario. 

63 The keys are the column names, and the values are the IDs of the associated PropertyValues. 

64 """ 

65 data_column: DataColumn 

66 data = { 

67 data_column.name: { 

68 "property": data_column.property_value.pk, 

69 "units": data_column.property_value.property.unit 

70 } for data_column in queryset.all() if data_column.property_value is not None 

71 } 

72 return data