feat: refactor get_asset_batch to async and replace simulation endpoint with test-asset-batch POST route

src/aeros_equipment/router.py

@@ -1,9 +1,9 @@
 
-from src.aeros_equipment.schema import EquipmentSimulation
 from src.aeros_equipment.service import get_asset_batch
-from typing import List, Optional
+from typing import List
 
 from fastapi import APIRouter, HTTPException, status
+from pydantic import BaseModel
 from sqlalchemy import select
 
 from src.aeros_equipment.model import AerosEquipmentCustomParameterData
@@ -11,7 +11,6 @@ from src.auth.service import CurrentUser
 from src.database.core import DbSession
 from src.database.service import CommonParameters
 from src.models import StandardResponse
-from src.config import RELIABILITY_SERVICE_API
 from .schema import CustomParameter, EquipmentPagination
 from .service import save_default_equipment, get_all
 
@@ -59,14 +58,19 @@ async def get_parameter_oreda(
         "message": "Success"
     }
 
-@router.get("/set_equipment_simulation", response_model=StandardResponse[object])
-async def set_equipment_simulation(
-    db_session: DbSession, location_tags: EquipmentSimulation
+class AssetBatchRequest(BaseModel):
+    repairable_tags: List[str]
+    non_repairable_tags: List[str]
+
+
+@router.post("/test-asset-batch", response_model=StandardResponse[dict])
+async def test_asset_batch(
+    db_session: DbSession, payload: AssetBatchRequest
 ):
-    equipments  = get_asset_batch(
-        location_tags,
-        location_tags,
-        RELIABILITY_SERVICE_API
+    equipments = await get_asset_batch(
+        payload.repairable_tags,
+        payload.non_repairable_tags,
+        db_session,
     )
 
     return {"data": equipments, "status": "success", "message": "Success"}

src/aeros_equipment/service.py

@@ -225,7 +225,7 @@ async def save_default_equipment(*, db_session: DbSession, project_name: str):
         )
 
 
-def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
+async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
                     db_session: DbSession,
                     base_url: str = RELIABILITY_SERVICE_API,
                    timeout: int = 600)-> dict:
@@ -255,11 +255,12 @@ def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
             WHERE location_tag = ANY(:location_tags)
         """)
         
-        repairable_result = db_session.execute(repairable_query, {"location_tags": location_tags})
+        repairable_result = await db_session.execute(repairable_query, {"location_tags": location_tags})
         
         for row in repairable_result:
             location_tag = row[0]
-            parameters = json.loads(row[1]) if row[1] else {}
+            raw_parameters = row[1]
+            parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
             
             try:
                 distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=True)
@@ -282,11 +283,12 @@ def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
             WHERE location_tag = ANY(:nr_location_tags)
         """)
         
-        non_repairable_result = db_session.execute(non_repairable_query, {"nr_location_tags": nr_location_tags})
+        non_repairable_result = await db_session.execute(non_repairable_query, {"nr_location_tags": nr_location_tags})
         
         for row in non_repairable_result:
             location_tag = row[0]
-            parameters = json.loads(row[1]) if row[1] else {}
+            raw_parameters = row[1]
+            parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
             best_distribution = row[2]
             
             try:
@@ -439,7 +441,7 @@ async def update_equipment_for_simulation(
         t.location_tag for t in non_repairable.scalars().all()
     ]
 
-    reliability_data = get_asset_batch(
+    reliability_data = await get_asset_batch(
         reg_nodes,
         non_repairable_location_tags,
         db_session