Coverage for backend/django/core/auxiliary/viewsets/CompoundViewSet.py: 100%

1from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes 

2from pydantic import BaseModel 

3from rest_framework import status, viewsets 

4from rest_framework.response import Response 

5from ahuora_compounds.CompoundDB import db 

6from ahuora_compounds.packages.chemsep import ChemsepCompoundData 

7from core.auxiliary.enums.unitOpGraphics import ConType 

8from flowsheetInternals.unitops.models.SimulationObject import SimulationObject 

9 

10 

11class CompoundListItem(BaseModel): 

12 name: str 

13 formula: str | None = None 

14 

15class CompoundViewSet(viewsets.ViewSet): 

16 """Viewset for retrieving compounds and their physical properties.""" 

17 

18 @extend_schema(responses={status.HTTP_200_OK: CompoundListItem}, parameters=[ 

19 OpenApiParameter(name="simulation_object_id", required=True, type=OpenApiTypes.INT) 

20 ],) 

21 def list(self, request): 

22 """List all compounds.""" 

23 simulation_object_id = request.query_params.get("simulation_object_id") 

24 simulation_object = SimulationObject.objects.get(id=simulation_object_id) 

25 property_package_key = simulation_object.propertyPackageType 

26 

27 def get_compound_item(compound_name: str) -> CompoundListItem: 

28 compound_info = ChemsepCompoundData.get(compound_name, None) 

29 

30 return CompoundListItem( 

31 name=compound_name, 

32 formula=compound_info.StructureFormula.value if compound_info else None, 

33 ) 

34 

35 compounds = [ 

36 get_compound_item(compound_name).model_dump() 

37 for compound_name in 

38 db.get_supported_compounds([property_package_key]) 

39 ] 

40 

41 return Response(compounds, status=status.HTTP_200_OK)