Coverage for backend/django/PinchAnalysis/views/HenNodeViewSet.py: 81%

1from PinchAnalysis.models.HenNode import HenNode 

2from django.db.models import Q 

3from django.db import transaction 

4from rest_framework import serializers 

5from PinchAnalysis.views.henNodeHelpers import group_stream_by_unitop_type, is_hot, is_cold, set_hennode_connections 

6from core.viewset import ModelViewSet 

7from rest_framework.decorators import action 

8from rest_framework.response import Response 

9from PinchAnalysis.models.InputModels import StreamDataEntry 

10from drf_spectacular.utils import extend_schema 

11import traceback 

12 

13 

14class BulkCreateHenNodeInputSerializer(serializers.Serializer): 

15 stream_ids = serializers.ListField( 

16 child=serializers.IntegerField(), 

17 required=False 

18 ) 

19 

20class HenNodeSerializer(serializers.ModelSerializer): 

21 class Meta: 

22 model = HenNode 

23 fields = "__all__" 

24 

25 

26class HenNodeViewSet(ModelViewSet): 

27 serializer_class = HenNodeSerializer 

28 

29 def get_queryset(self): 

30 return HenNode.objects.all() 

31 

32 def error_response(self, e): 

33 tb_info = traceback.format_exc() 

34 error_message = str(e) 

35 response_data = { 

36 'status': 'error', 

37 'message': error_message, 

38 'traceback': tb_info 

39 } 

40 return Response(response_data, status=400)