Coverage for backend/django/flowsheetInternals/unitops/viewsets/SimulationObjectViewSet.py: 82%

237 statements  

« prev     ^ index     » next       coverage.py v7.10.7, created at 2026-06-23 21:51 +0000

1import json 

2 

3from core.auxiliary.models.PropertyValue import PropertyValue 

4from core.auxiliary.enums.unitOpData import SimulationObjectClass 

5from core.viewset import ModelViewSet 

6from rest_framework.response import Response 

7from flowsheetInternals.unitops.viewsets.DuplicateSimulationObject import DuplicateSimulationObject 

8from flowsheetInternals.unitops.models.compound_propogation import update_compounds_on_set 

9from flowsheetInternals.unitops.models.flow_tracking import track_downstream_stream_flow 

10from flowsheetInternals.unitops.logic.insert_translator_block import insert_translator_block 

11from flowsheetInternals.unitops.serializers.SimulationObjectSerializer import SimulationObjectRetrieveSerializer 

12from flowsheetInternals.unitops.models import SimulationObject 

13from flowsheetInternals.unitops.models.delete_factory import DeleteFactory 

14from flowsheetInternals.unitops.serializers import * 

15from core.auxiliary.models.PropertySet import PropertySet 

16from core.auxiliary.models.PropertyInfo import PropertyInfo 

17from flowsheetInternals.unitops.models.summary_table_factory import get_composition_summary_table_data, get_stream_summary_table_data, get_unitops_summary_table_data 

18from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes 

19from rest_framework import serializers 

20from rest_framework.decorators import action 

21from rest_framework import status 

22import traceback 

23from django.db import transaction 

24from django.db.models import Prefetch 

25from common.config_types import PropertyType as PropertyTypeObj 

26from flowsheetInternals.formula_templates.add_template import add_predefined_template 

27from flowsheetInternals.unitops.methods.add_expression import add_expression 

28from flowsheetInternals.graphicData.models.groupingModel import Grouping 

29 

30class UpdateCompoundSerializer(serializers.Serializer): 

31 """Payload schema for updating compound selections on a stream.""" 

32 

33 compounds = serializers.ListField(child=serializers.CharField()) 

34 simulationObject = serializers.IntegerField() 

35 

36class AddPropertyTemplateSerializer(serializers.Serializer): 

37 """Payload schema for adding a property template by key to a simulation object.""" 

38 

39 templateKey = serializers.CharField() 

40 

41class DuplicateSimulationObjectSerializer(serializers.Serializer): 

42 """Payload schema for duplicating one or more simulation objects.""" 

43 

44 objectIDs = serializers.ListField(child=serializers.IntegerField()) 

45 x = serializers.FloatField() 

46 y = serializers.FloatField() 

47 

48class AddPortSerializer(serializers.Serializer): 

49 """Payload schema for attaching a new port to a simulation object.""" 

50 

51 simulationObjectId = serializers.IntegerField() 

52 key = serializers.CharField() 

53 stream = serializers.IntegerField(required=False, allow_null=True) #optional if existing stream isnt passed in 

54 

55class MergeDecisionNodesSerializer(serializers.Serializer): 

56 """Payload schema for merging two decision nodes.""" 

57 

58 decisionNodeActive = serializers.IntegerField(required=True) 

59 decisionNodeOver = serializers.IntegerField(required=True) 

60 

61class TrackedStreamFlowSerializer(serializers.Serializer): 

62 """Response schema for stream-flow highlighting.""" 

63 

64 sourceStreamId = serializers.IntegerField() 

65 streamIds = serializers.ListField(child=serializers.IntegerField()) 

66 

67class RestoreObjectsSerializer(serializers.Serializer): 

68 """Payload schema for restoring previously deleted simulation objects.""" 

69 

70 object_ids = serializers.ListField( 

71 child=serializers.IntegerField(), 

72 help_text="Array of object IDs to restore", 

73 min_length=1 

74 ) 

75 

76 def validate(self, data): 

77 """Ensure we have at least one valid ID""" 

78 if not data.get('object_ids'): 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true

79 raise serializers.ValidationError("At least one object ID must be provided") 

80 return data 

81 

82class SimulationObjectViewSet(ModelViewSet): 

83 """API endpoints for interacting with flowsheet simulation objects.""" 

84 

85 serializer_class = SimulationObjectSerializer 

86 

87 def get_queryset(self): 

88 """Return simulation objects with related property metadata eagerly loaded.""" 

89 # Prefetch related entities to avoid N+1 queries when serializing. 

90 queryset = SimulationObject.objects.all().select_related( 

91 "properties", 

92 "grouping" 

93 ).prefetch_related( 

94 Prefetch( 

95 "properties__ContainedProperties", 

96 queryset=PropertyInfo.objects 

97 .select_related("recycleConnection") 

98 .prefetch_related( 

99 Prefetch("values", queryset=PropertyValue.objects 

100 .select_related("controlManipulated", "controlSetPoint") 

101 .prefetch_related("indexedItems")) 

102 ) 

103 ), 

104 "connectedPorts__unitOp" 

105 ) 

106 

107 return queryset 

108 

109 def get_serializer_class(self): 

110 """Use the retrieve serializer when properties must be serialized.""" 

111 if self.action in ["list", "retrieve"]: 

112 # For the list and retrieve actions, include the properties of the 

113 # simulation object so canvas consumers can render from one payload. 

114 return SimulationObjectRetrieveSerializer 

115 return SimulationObjectSerializer 

116 

117 @extend_schema( 

118 parameters=[ 

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

120 ] 

121 ) 

122 def list(self, request): 

123 """List simulation objects filtered by the provided flowsheet ID.""" 

124 return super().list(request) 

125 

126 def destroy(self, request, *args, **kwargs): 

127 """Delete the simulation object via the shared DeleteFactory.""" 

128 instance = self.get_object() 

129 DeleteFactory.delete_object(instance) 

130 

131 return Response(status=status.HTTP_204_NO_CONTENT) 

132 #maybe do local storage and its request it in here 

133 

134 @extend_schema(responses=TrackedStreamFlowSerializer) 

135 @action(detail=True, methods=["get"], url_path="tracked-stream-flow") 

136 def tracked_stream_flow(self, request, pk=None): 

137 """Return the downstream streams connected to this stream.""" 

138 simulation_object = self.get_object() 

139 stream_types = { 

140 SimulationObjectClass.Stream, 

141 SimulationObjectClass.HumidAirStream, 

142 } 

143 

144 if simulation_object.objectType not in stream_types: 

145 return Response( 

146 {"detail": "Track stream is only available for streams."}, 

147 status=status.HTTP_400_BAD_REQUEST, 

148 ) 

149 

150 _unit_ops, streams = track_downstream_stream_flow(simulation_object) 

151 stream_ids = sorted(stream.id for stream in streams) 

152 

153 return Response( 

154 { 

155 "sourceStreamId": simulation_object.id, 

156 "streamIds": stream_ids, 

157 }, 

158 status=status.HTTP_200_OK, 

159 ) 

160 

161 

162 def get_summary_queryset(self): 

163 queryset = SimulationObject.objects.all().select_related( 

164 "properties", 

165 ).prefetch_related( 

166 Prefetch( 

167 "properties__ContainedProperties", 

168 queryset=PropertyInfo.objects 

169 .prefetch_related( 

170 Prefetch("values", queryset=PropertyValue.objects 

171 .prefetch_related("indexedItems")) 

172 ) 

173 ) 

174 ) 

175 

176 return queryset 

177 

178 @extend_schema( 

179 parameters=[ 

180 OpenApiParameter(name="unit_map", type=OpenApiTypes.STR, required=True), 

181 OpenApiParameter(name="groups", type=OpenApiTypes.ANY, required=True) 

182 ], 

183 responses=serializers.DictField() 

184 ) 

185 @action(detail=False, methods=['get'], url_path='streams-summary') 

186 def summary_table_streams(self, request): 

187 unit_map = self.request.query_params.get("unit_map") 

188 groups = self.request.query_params.get("groups") 

189 

190 results = {} 

191 groups = Grouping.objects.filter(id__in=json.loads(groups)) 

192 for current_group in groups: 

193 query_set = self.get_summary_queryset().filter(graphicObject__in=current_group.graphicObjects.all()) 

194 

195 data = get_stream_summary_table_data(query_set, json.loads(unit_map)) 

196 

197 group_name = current_group.simulationObject.componentName 

198 results[group_name] = data 

199 

200 return Response(results, status=200) 

201 

202 @extend_schema( 

203 parameters=[ 

204 OpenApiParameter(name="unit_map", type=OpenApiTypes.STR, required=True), 

205 OpenApiParameter(name="groups", type=OpenApiTypes.ANY, required=True) 

206 ], 

207 responses=serializers.DictField() 

208 ) 

209 @action(detail=False, methods=['get'], url_path='unitops-summary') 

210 def summary_table_unitops(self, request): 

211 unit_map = self.request.query_params.get("unit_map") 

212 groups = self.request.query_params.get("groups") 

213 

214 results = {} 

215 groups = Grouping.objects.filter(id__in=json.loads(groups)) 

216 for current_group in groups: 

217 query_set = self.get_summary_queryset().filter(graphicObject__in=current_group.graphicObjects.all()) 

218 data = get_unitops_summary_table_data(query_set, json.loads(unit_map)) 

219 

220 group_name = current_group.simulationObject.componentName 

221 results[group_name] = data 

222 

223 return Response(results, status=200) 

224 

225 @extend_schema( 

226 parameters=[ 

227 OpenApiParameter(name="compound_mode", type=OpenApiTypes.STR, required=True), 

228 OpenApiParameter(name="measure_type", type=OpenApiTypes.STR, required=True), 

229 OpenApiParameter(name="groups", type=OpenApiTypes.ANY, required=True) 

230 ], 

231 responses=serializers.DictField() 

232 ) 

233 @action(detail=False, methods=['get'], url_path='compounds-summary') 

234 def summary_table_compounds(self, request): 

235 compound_mode = self.request.query_params.get("compound_mode") 

236 measure_type = self.request.query_params.get("measure_type") 

237 groups = self.request.query_params.get("groups") 

238 

239 results = {} 

240 groups = Grouping.objects.filter(id__in=json.loads(groups)) 

241 for current_group in groups: 

242 query_set = self.get_summary_queryset().filter(graphicObject__in=current_group.graphicObjects.all()) 

243 data = get_composition_summary_table_data(queryset=query_set, target_compound_mode=compound_mode, measure_type=measure_type) 

244 

245 group_name = current_group.simulationObject.componentName 

246 results[group_name] = data 

247 

248 return Response(results, status=200) 

249 

250 @extend_schema(request=UpdateCompoundSerializer, responses=None) 

251 @action(detail=False, methods=['post'], url_path='update-compounds') 

252 def update_compounds(self, request): 

253 """Update the selected compounds for the specified simulation object.""" 

254 try: 

255 serializer = UpdateCompoundSerializer(data=request.data) 

256 serializer.is_valid(raise_exception=True) 

257 validated_data = serializer.validated_data 

258 

259 simulation_object = SimulationObject.objects.get( 

260 pk=validated_data.get("simulationObject") 

261 ) 

262 

263 expected_compounds = validated_data.get("compounds") 

264 update_compounds_on_set(simulation_object, expected_compounds) 

265 

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

267 

268 except (SimulationObject.DoesNotExist, PropertySet.DoesNotExist) as e: 

269 return Response({'status': 'error', 'message': str(e)}, status=404) 

270 except Exception as e: 

271 print(traceback.format_exc()) 

272 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400) 

273 

274 

275 

276 @extend_schema(request=AddPortSerializer, responses=None) 

277 @action(detail=False, methods=['post'], url_path='add-port') 

278 def add_port(self, request): 

279 """Add a port to a simulation object, optionally reusing an existing stream.""" 

280 try: 

281 serializer = AddPortSerializer(data=request.data) 

282 serializer.is_valid(raise_exception=True) 

283 validated_data = serializer.validated_data 

284 

285 simulation_object = SimulationObject.objects.get( 

286 pk=validated_data.get("simulationObjectId") 

287 ) 

288 

289 key = validated_data.get("key") 

290 

291 streamId = validated_data.get("stream") 

292 existing_stream = None 

293 

294 if streamId is not None: 294 ↛ 295line 294 didn't jump to line 295 because the condition on line 294 was never true

295 existing_stream = SimulationObject.objects.get(pk=streamId) 

296 # Reuse the existing stream to keep compound propagation in sync. 

297 simulation_object.add_port(key, existing_stream) 

298 else: 

299 simulation_object.add_port(key) 

300 

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

302 

303 except (SimulationObject.DoesNotExist, PropertySet.DoesNotExist) as e: 

304 return Response({'status': 'error', 'message': str(e)}, status=404) 

305 except Exception as e: 

306 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400) 

307 

308 

309 

310 @extend_schema(request=MergeDecisionNodesSerializer, responses=None) 

311 @action(detail=False, methods=['post'], url_path='merge-decision-nodes') 

312 def merge_decision_nodes(self, request): 

313 try: 

314 serializer = MergeDecisionNodesSerializer(data=request.data) 

315 serializer.is_valid(raise_exception=True) 

316 validated_data = serializer.validated_data 

317 

318 decision_node_active = SimulationObject.objects.get( 

319 pk=validated_data.get("decisionNodeActive") 

320 ) 

321 

322 decision_node_over = SimulationObject.objects.get( 

323 pk=validated_data.get("decisionNodeOver") 

324 ) 

325 

326 decision_node_active.merge_decision_nodes(decision_node_active, decision_node_over) 

327 

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

329 

330 except (SimulationObject.DoesNotExist, PropertySet.DoesNotExist) as e: 

331 return Response({'status': 'error', 'message': str(e)}, status=404) 

332 except Exception as e: 

333 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400) 

334 

335 

336 

337 @extend_schema(request=None,responses=int) 

338 @action(detail=True, methods=['POST']) 

339 def add_expression(self, request, pk=None): 

340 """Create a blank expression property on the simulation object.""" 

341 if pk is None: 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true

342 return Response({'status': 'error', 'message': 'id is required'}, status=400) 

343 

344 simulation_object = SimulationObject.objects.get(pk=pk) 

345 prop_info = add_expression(simulation_object) 

346 

347 return Response(prop_info.id, status=200) 

348 

349 @extend_schema(request=AddPropertyTemplateSerializer) 

350 @action(detail=True, methods=['POST']) 

351 def add_custom_property_template(self, request, pk=None): 

352 """Attach a custom propertytemplate to the simulation object using the supplied key.""" 

353 if pk is None: 353 ↛ 354line 353 didn't jump to line 354 because the condition on line 353 was never true

354 return Response({'status': 'error', 'message': 'id is required'}, status=400) 

355 

356 simulation_object = SimulationObject.objects.prefetch_related( 

357 "properties" 

358 ).get(pk=pk) 

359 serializer = AddPropertyTemplateSerializer(data=request.data) 

360 serializer.is_valid(raise_exception=True) 

361 template_key = serializer.validated_data.get("templateKey") 

362 

363 try: 

364 add_predefined_template(simulation_object, template_key) 

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

366 except ValueError as e: 

367 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400) 

368 

369 

370 @extend_schema(request=RestoreObjectsSerializer, responses=None) 

371 @action(detail=False, methods=['post']) 

372 def restore(self, request): 

373 """Restore previously deleted simulation objects.""" 

374 try: 

375 # Handle both 'ids' and 'object_ids' in request data 

376 if 'ids' in request.data: 376 ↛ 377line 376 didn't jump to line 377 because the condition on line 376 was never true

377 request_data = {'object_ids': request.data['ids']} 

378 else: 

379 request_data = request.data 

380 

381 serializer = RestoreObjectsSerializer(data=request_data) 

382 serializer.is_valid(raise_exception=True) 

383 object_ids = serializer.validated_data['object_ids'] 

384 except Exception as e: 

385 return Response( 

386 { 

387 'error': 'Validation error', 

388 'details': str(e), 

389 'traceback': traceback.format_exc(), 

390 'received_data': request.data 

391 }, 

392 status=status.HTTP_400_BAD_REQUEST 

393 ) 

394 

395 if not object_ids: 395 ↛ 396line 395 didn't jump to line 396 because the condition on line 395 was never true

396 return Response( 

397 {'error': 'No objectIds provided'}, 

398 status=status.HTTP_400_BAD_REQUEST 

399 ) 

400 

401 DeleteFactory._restore_object_ids(object_ids) 

402 return Response(status=status.HTTP_200_OK) 

403 

404 

405 

406 @extend_schema(request=DuplicateSimulationObjectSerializer, responses=None) 

407 @action(detail=False, methods=['post'], url_path='duplicate-simulation-object') 

408 def duplicate_simulation_object(self, request): 

409 """Duplicate simulation objects and reposition them at the supplied coordinates.""" 

410 try: 

411 serializer = DuplicateSimulationObjectSerializer(data=request.data) 

412 serializer.is_valid(raise_exception=True) 

413 validated_data = serializer.validated_data 

414 duplicator = DuplicateSimulationObject() 

415 flowsheet = self.request.query_params.get("flowsheet") 

416 duplicator.handle_duplication_request( 

417 flowsheet, 

418 validated_data 

419 ) 

420 return Response(None, status=200) 

421 except SimulationObject.DoesNotExist as e: 

422 return Response({'status': 'error', 'message': str(e)}, status=400) 

423 except Exception as e: 

424 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400) 

425 

426 

427 @extend_schema(request=None, responses=None) 

428 @action(detail=True, methods=['POST'], url_path='insert-translator-block') 

429 def insert_translator_block(self, request, pk=None): 

430 """ 

431 Inserts a translator block into the flowsheet. 

432 """ 

433 if pk is None: 433 ↛ 434line 433 didn't jump to line 434 because the condition on line 433 was never true

434 return Response({'status': 'error', 'message': 'id is required'}, status=400) 

435 

436 try: 

437 with transaction.atomic(): 

438 stream = ( 

439 SimulationObject.objects 

440 .select_related("flowsheet") 

441 .prefetch_related( 

442 "graphicObject", 

443 "connectedPorts" 

444 ) 

445 .get(pk=pk) 

446 ) 

447 insert_translator_block(stream) 

448 

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

450 except Exception as e: 

451 return Response({'status': 'error', 'message': str(e),'traceback': traceback.format_exc()}, status=400)