Coverage for backend/core/auxiliary/viewsets/MLColumnMappingsViewSet.py: 67%

1from core.viewset import ModelViewSet 

2from core.auxiliary.models.MLColumnMapping import MLColumnMapping 

3from core.auxiliary.models.Flowsheet import Flowsheet 

4from rest_framework import serializers 

5from core.auxiliary.serializers.MLColumnMappingsSerializer import MLColumnMappingsSerializer 

6from rest_framework.response import Response 

7from drf_spectacular.utils import extend_schema 

8from rest_framework.decorators import action 

9from core.auxiliary.serializers.PropertyInfoSerializer import PropertyInfoSerializer 

10from drf_spectacular.utils import extend_schema 

11from rest_framework.exceptions import ValidationError 

12from core.auxiliary.models.MLWizard import create_column_mapping 

13from django.db import transaction 

14 

15class ColumnMappingSerializer(serializers.Serializer): 

16 portIndex = serializers.IntegerField() 

17 propertyKey = serializers.CharField() 

18 column = serializers.CharField() 

19 

20class BulkCreateColumnMappingSerializer(serializers.Serializer): 

21 model = serializers.IntegerField() 

22 inlet_mappings = serializers.ListField(child=ColumnMappingSerializer()) 

23 outlet_mappings = serializers.ListField(child=ColumnMappingSerializer()) 

24class PropertyInfoIdsSerializer(serializers.Serializer): 

25 inlets = serializers.ListField(child=PropertyInfoSerializer()) 

26 outlets = serializers.ListField(child=PropertyInfoSerializer()) 

27 

28 

29class MLColumnMappingsViewSet(ModelViewSet): 

30 serializer_class = MLColumnMappingsSerializer 

31 

32 def get_queryset(self): 

33 return MLColumnMapping.objects.all() 

34 

35 @transaction.atomic 

36 @extend_schema(request=BulkCreateColumnMappingSerializer, responses=None) 

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

38 def bulk_create_column_mapping(self, request): 

39 serializer = BulkCreateColumnMappingSerializer(data=request.data) 

40 serializer.is_valid(raise_exception=True) 

41 validated_data = serializer.validated_data 

42 model = validated_data.get("model") 

43 inlet_mappings = validated_data.get("inlet_mappings") 

44 outlet_mappings = validated_data.get("outlet_mappings") 

45 flowsheet = request.query_params.get("flowsheet") 

46 flowsheet = Flowsheet.objects.get(pk=flowsheet) 

47 try: 

48 create_column_mapping(flowsheet, model, inlet_mappings, outlet_mappings) 

49 except Exception as e: 

50 print(ValidationError(str(e))) 

51 return Response({'status': 'error'}, status=400) 

52 

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