feat: add execute_update_node to worker activity registration

run_worker.py

@@ -3,7 +3,7 @@ import os
 from temporalio.client import Client
 from temporalio.worker import Worker
 
-from temporal.activity import calculate_plant_eaf_activity, execute_simulation_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity, call_callback_ahm, acquire_lock, release_lock
+from temporal.activity import calculate_plant_eaf_activity, execute_simulation_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity, call_callback_ahm, acquire_lock, release_lock,execute_update_node
 from temporal.workflow import SimulationWorkflow
 
 TEMPORAL_URL = os.environ.get("TEMPORAL_URL", "http://192.168.1.86:7233")
@@ -23,7 +23,8 @@ async def main():
                 update_contribution_bulk_mappings_activity,
                 call_callback_ahm,
                 acquire_lock,
-                release_lock
+                release_lock,
+                execute_update_node
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=12

src/aeros_equipment/service.py

@@ -159,7 +159,7 @@ async def update_node(
     *, db_session: DbSession, equipment_nodes: List[dict], project_name: str
 ):
     updateNodeReq = {"projectName": project_name, "regNodeInputs": equipment_nodes}
-
+    log.info("Aeros UpdateEquipmentDistributions request: %s", updateNodeReq)
 
     try:
         response = await aeros_post(
@@ -167,14 +167,15 @@ async def update_node(
             json=updateNodeReq,
             headers={"Content-Type": "application/json"},
         )
-
+        response.raise_for_status()
         result = response.json()
-
+        log.info("Aeros UpdateEquipmentDistributions response: %s", result)
         return result
 
     except Exception as e:
+        log.error("Failed to update equipment nodes in Aeros: %s", str(e))
         raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Aeros update failed: {str(e)}"
         )
 
 
@@ -259,21 +260,26 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
         
         for row in repairable_result:
             location_tag = row[0]
+            if not location_tag:
+                continue
+            
+            # Normalize key
+            tag_key = location_tag.strip().upper()
             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)
                 
-                results[location_tag]["cmDisType"] = "Normal"
+                results[tag_key]["cmDisType"] = "Normal"
-                results[location_tag]["cmDisP1"] = 6
+                results[tag_key]["cmDisP1"] = 6
-                results[location_tag]["cmDisP2"] = 3
+                results[tag_key]["cmDisP2"] = 3
-                results[location_tag]["relDisType"] = distribution
+                results[tag_key]["relDisType"] = distribution
-                results[location_tag]["relDisP1"] = reldisp1
+                results[tag_key]["relDisP1"] = reldisp1
-                results[location_tag]["relDisP2"] = reldisp2
+                results[tag_key]["relDisP2"] = reldisp2
-                results[location_tag]["parameters"] = parameters
+                results[tag_key]["parameters"] = parameters
             except Exception as e:
-                raise Exception(f"Error processing repairable item {location_tag}: {e}")
+                log.error(f"Error processing repairable item {location_tag}: {e}")
 
     # Query non-repairable equipment data
     if nr_location_tags:
@@ -287,6 +293,11 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
         
         for row in non_repairable_result:
             location_tag = row[0]
+            if not location_tag:
+                continue
+                
+            # Normalize key
+            tag_key = location_tag.strip().upper()
             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]
@@ -294,15 +305,15 @@ async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
             try:
                 distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=False, best_distribution=best_distribution)
                 
-                results[location_tag]["cmDisType"] = "Normal"
+                results[tag_key]["cmDisType"] = "Normal"
-                results[location_tag]["cmDisP1"] = 6
+                results[tag_key]["cmDisP1"] = 6
-                results[location_tag]["cmDisP2"] = 3
+                results[tag_key]["cmDisP2"] = 3
-                results[location_tag]["relDisType"] = distribution
+                results[tag_key]["relDisType"] = distribution
-                results[location_tag]["relDisP1"] = reldisp1
+                results[tag_key]["relDisP1"] = reldisp1
-                results[location_tag]["relDisP2"] = reldisp2
+                results[tag_key]["relDisP2"] = reldisp2
-                results[location_tag]["parameters"] = parameters
+                results[tag_key]["parameters"] = parameters
             except Exception as e:
-                raise Exception(f"Error processing non-repairable item {location_tag}: {e}")
+                log.error(f"Error processing non-repairable item {location_tag}: {e}")
 
     return results
 
@@ -366,7 +377,7 @@ def get_distribution_from_db(parameters: dict, is_repairable: bool = True, best_
             return "Normal", mu, sigma
     
 async def update_oh_interval_offset(*, aeros_db_session, overhaul_offset, overhaul_interval, project_name):
-    query = text("""
+    stmt = text("""
         UPDATE public."RegularNodes" rn
         SET "OHInterval" = :new_oh_interval,
             "OHOffset" = :new_oh_offset
@@ -377,13 +388,14 @@ async def update_oh_interval_offset(*, aeros_db_session, overhaul_offset, overha
         AND p."ProjectName" = :project_name
     """)
 
-    await aeros_db_session.execute(query, {
+    result = await aeros_db_session.execute(stmt, {
         "new_oh_interval": overhaul_interval,
         "new_oh_offset": overhaul_offset,
         "project_name": project_name
     })
     
     await aeros_db_session.commit()
+    log.info("Updated OH Interval/Offset for project %s. Rows affected: %s", project_name, result.rowcount)
 
 
 async def update_equipment_for_simulation(
@@ -528,7 +540,15 @@ async def update_equipment_for_simulation(
                         eq["relDisP1"] = 1.0
                         eq["relDisP2"] = MTBF
 
+                eq["relDisP3"] = 0
+                eq["relDisUnitCode"] = "UHour"
+                eq["cmDisP2"] = 0
+                eq["cmDisP3"] = 0
+                eq["cmDisUnitCode"] = "UHour"
+                eq["ohDisType"] = "Fixed"
                 eq["ohDisP1"] = overhaul_duration
+                eq["ohDisP2"] = 0
+                eq["ohDisP3"] = 0
                 eq["ohDisUnitCode"] = "UHour"
 
                 reqNodeInputs.append(eq)
@@ -547,24 +567,48 @@ async def update_equipment_for_simulation(
             if eq_name_raw in results:
                 continue
 
-            reliability = reliability_data.get(eq_name_raw, {})
+            # Lookup reliability with normalized key
+            reliability = reliability_data.get(eq_name, {})
 
+            if reliability:
+                log.info("Found reliability data for %s", eq_name_raw)
                 eq["cmDisType"] = reliability.get("cmDisType", "Fixed")
                 eq["cmDisP1"] = reliability.get("cmDisP1", 0)
                 eq["cmDisP2"] = reliability.get("cmDisP2", 0)
+                eq["cmDisP3"] = 0
+                eq["cmDisUnitCode"] = "UHour"
                 eq["relDisType"] = reliability.get("relDisType", "Fixed")
                 eq["relDisP1"] = reliability.get("relDisP1", 0)
                 eq["relDisP2"] = reliability.get("relDisP2", 0)
+                eq["relDisP3"] = 0
+                eq["relDisUnitCode"] = "UHour"
+            else:
+                log.debug("No reliability data for %s, keeping existing or using defaults", eq_name_raw)
+                # Ensure existing unit codes and types are at least present if not already
+                eq.setdefault("cmDisType", "Fixed")
+                eq.setdefault("cmDisUnitCode", "UHour")
+                eq.setdefault("relDisType", "Fixed")
+                eq.setdefault("relDisUnitCode", "UHour")
+                eq.setdefault("cmDisP1", 0)
+                eq.setdefault("cmDisP2", 0)
+                eq.setdefault("cmDisP3", 0)
+                eq.setdefault("relDisP1", 0)
+                eq.setdefault("relDisP2", 0)
+                eq.setdefault("relDisP3", 0)
+
+            eq["ohDisType"] = "Fixed"
             eq["ohDisP1"] = overhaul_duration
+            eq["ohDisP2"] = 0
+            eq["ohDisP3"] = 0
             eq["ohDisUnitCode"] = "UHour"
 
             reqNodeInputs.append(eq)
             results[eq_name_raw] = {
-                "mttr": eq["cmDisP1"],
+                "mttr": eq.get("cmDisP1", 0),
-                "distribution": eq["relDisType"],
+                "distribution": eq.get("relDisType", "Fixed"),
-                "beta": eq["relDisP1"],
+                "beta": eq.get("relDisP1", 0),
-                "eta": eq["relDisP2"],
+                "eta": eq.get("relDisP2", 0),
-                "parameters": eq.get("parameters", {}),
+                "parameters": reliability.get("parameters", {}),
                 "oh_duration": overhaul_duration
             }
 
@@ -574,13 +618,13 @@ async def update_equipment_for_simulation(
             )
             reqNodeInputs.append(eq)
 
-    await update_node(
+    # await update_node(
-        db_session=db_session,
+    #     db_session=db_session,
-        equipment_nodes=reqNodeInputs,
+    #     equipment_nodes=reqNodeInputs,
-        project_name=project_name
+    #     project_name=project_name
-    )
+    # )
 
-    return results
+    return reqNodeInputs    
 
 # Optimized individual fetch functions
 async def get_equipment_mttr(*, location_tag: str, client: httpx.AsyncClient) -> float:

src/aeros_simulation/router.py

@@ -648,6 +648,7 @@ airflow_router = APIRouter()
 @airflow_router.post("/calculate_eaf_contribution", response_model=StandardResponse[dict])
 async def calculate_contribution(
     db_session: DbSession,
+    aeros_db_session: CollectorDbSession,
     simulation_in: SimulationInput,
     batch_num: int = Query(0, ge=0)
 ):
@@ -700,7 +701,15 @@ async def calculate_contribution(
             }
 
             results = await update_equipment_for_simulation(
-                db_session=db_session, project_name=project.project_name, schematic_name=simulation_in.SchematicName, custom_input=custom_input
+                db_session=db_session,
+                aeros_db_session=aeros_db_session,
+                project_name=project.project_name,
+                simulation_duration=simulation_in.SimDuration,
+                overhaul_duration=simulation_in.OverhaulDuration,
+                overhaul_interval=simulation_in.OverhaulInterval,
+                offset=simulation_in.OffSet,
+                schematic_name=simulation_in.SchematicName,
+                custom_input=custom_input
             )
 
             # await update_simulation(

temporal/activity.py

@@ -52,6 +52,16 @@ async def execute_simulation_activity(params: dict):
             eq_update=params["eq_update"]
         )
 
+@activity.defn
+async def execute_update_node(params:dict):
+    from src.database.core import async_session
+    from src.aeros_equipment.service import update_node
+    async with async_session() as db_session:
+        return await update_node(
+            db_session=db_session,
+            project_name=params["projectName"],
+            equipment_nodes=params["equipment_nodes"]
+        )
 
 @activity.defn
 async def calculate_plant_eaf_activity(params: dict):

temporal/workflow.py

@@ -2,7 +2,7 @@ from datetime import timedelta
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from temporal.activity import calculate_plant_eaf_activity, execute_simulation_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity,call_callback_ahm, acquire_lock, release_lock
+    from temporal.activity import calculate_plant_eaf_activity, execute_simulation_activity, update_contribution_bulk_mappings_activity, update_equipment_for_simulation_activity,call_callback_ahm, acquire_lock, release_lock,execute_update_node
 
 
 
@@ -39,6 +39,12 @@ class SimulationWorkflow:
             start_to_close_timeout=timedelta(days=1)
         )
 
+        await workflow.execute_activity(
+            execute_update_node,
+            {"projectName": sim_data['projectName'], "equipment_nodes": results},
+            start_to_close_timeout=timedelta(days=1)
+        )
+
         # 2. Execute simulation
         await workflow.execute_activity(
             execute_simulation_activity,

tests/unit/test_update_equipment_simulation.py

@@ -0,0 +1,227 @@
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
+from src.aeros_equipment.service import update_equipment_for_simulation
+from src.aeros_equipment.model import AerosEquipment, AerosEquipmentGroup, ReliabilityPredictNonRepairable
+from src.aeros_simulation.model import AerosNode
+import logging
+
+@pytest.mark.asyncio
+async def test_update_equipment_for_simulation_basic():
+    # Mock parameters
+    db_session = AsyncMock()
+    aeros_db_session = AsyncMock()
+    project_name = "TestProject"
+    simulation_duration = 8760
+    overhaul_duration = 72
+    overhaul_interval = 8000
+    offset = 0
+    schematic_name = "TestSchematic"
+    
+    # Mock get_aeros_schematic_by_name
+    mock_schematic = MagicMock()
+    mock_schematic.id = 1
+    
+    # Mock db_session.execute results
+    mock_eq = MagicMock()
+    mock_eq.location_tag = "EQ001"
+    mock_eq_result = MagicMock()
+    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
+    
+    mock_group = MagicMock()
+    mock_group.group_name = "GROUP001"
+    mock_group_result = MagicMock()
+    mock_group_result.scalars.return_value.all.return_value = [mock_group]
+    
+    mock_nr_result = MagicMock()
+    mock_nr_result.scalars.return_value.all.return_value = []
+
+    db_session.execute.side_effect = [
+        mock_eq_result,    # for equipments
+        mock_group_result, # for rbd_group
+        mock_nr_result     # for non_repairable
+    ]
+    
+    # Mock nodes data
+    mock_nodes_data = [
+        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0},
+        {"equipmentName": "GROUP001", "relDisType": "Fixed", "relDisP1": 0}
+    ]
+    
+    # Mock reliability data
+    mock_reliability_data = {
+        "EQ001": {
+            "relDisType": "Weibull2",
+            "relDisP1": 1.5,
+            "relDisP2": 10000,
+            "cmDisType": "Fixed",
+            "cmDisP1": 5,
+            "cmDisP2": 0
+        }
+    }
+
+    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
+         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
+         patch("src.aeros_equipment.service.get_asset_batch", return_value=mock_reliability_data), \
+         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
+         patch("src.aeros_equipment.service.update_node", return_value=None) as mock_update_node:
+        
+        # Run the function
+        results = await update_equipment_for_simulation(
+            db_session=db_session,
+            aeros_db_session=aeros_db_session,
+            project_name=project_name,
+            simulation_duration=simulation_duration,
+            overhaul_duration=overhaul_duration,
+            overhaul_interval=overhaul_interval,
+            offset=offset,
+            schematic_name=schematic_name
+        )
+        
+        # Assertions
+        assert "EQ001" in results
+        assert results["EQ001"]["distribution"] == "Weibull2"
+        assert results["EQ001"]["oh_duration"] == overhaul_duration
+        
+        # Check if update_node was called with correct arguments
+        called_args = mock_update_node.call_args[1]
+        equipment_nodes = called_args["equipment_nodes"]
+        
+        # Find EQ001 in equipment_nodes
+        eq001_node = next(node for node in equipment_nodes if node["equipmentName"] == "EQ001")
+        assert eq001_node["relDisType"] == "Weibull2"
+        assert eq001_node["ohDisP1"] == overhaul_duration
+        assert eq001_node["relDisUnitCode"] == "UHour"
+        assert eq001_node["ohDisUnitCode"] == "UHour"
+
+@pytest.mark.asyncio
+async def test_update_equipment_for_simulation_custom_input():
+    # Mock parameters
+    db_session = AsyncMock()
+    aeros_db_session = AsyncMock()
+    project_name = "TestProject"
+    simulation_duration = 8760
+    overhaul_duration = 72
+    overhaul_interval = 8000
+    offset = 0
+    schematic_name = "TestSchematic"
+    
+    custom_input = {
+        "EQ001": {
+            "mttr": 10,
+            "failure_rate": 0.5,
+            "distribution": "Exponential2"
+        }
+    }
+    
+    # Mock get_aeros_schematic_by_name
+    mock_schematic = MagicMock()
+    mock_schematic.id = 1
+    
+    # Mock db_session.execute
+    mock_eq = MagicMock()
+    mock_eq.location_tag = "EQ001"
+    mock_eq_result = MagicMock()
+    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
+    
+    mock_empty_result = MagicMock()
+    mock_empty_result.scalars.return_value.all.return_value = []
+
+    db_session.execute.side_effect = [
+        mock_eq_result,    # equipments
+        mock_empty_result, # rbd_group
+        mock_empty_result  # non_repairable
+    ]
+    
+    # Mock nodes data
+    mock_nodes_data = [
+        {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0}
+    ]
+    
+    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
+         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
+         patch("src.aeros_equipment.service.get_asset_batch", return_value={}), \
+         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
+         patch("src.aeros_equipment.service.update_node", return_value=None):
+        
+        # Run
+        results = await update_equipment_for_simulation(
+            db_session=db_session,
+            aeros_db_session=aeros_db_session,
+            project_name=project_name,
+            simulation_duration=simulation_duration,
+            overhaul_duration=overhaul_duration,
+            overhaul_interval=overhaul_interval,
+            offset=offset,
+            schematic_name=schematic_name,
+            custom_input=custom_input
+        )
+        
+        assert "EQ001" in results
+        assert results["EQ001"]["mttr"] == 10
+        assert results["EQ001"]["distribution"] == "Exponential2"
+
+@pytest.mark.asyncio
+async def test_update_equipment_for_simulation_missing_reliability():
+    # Mock parameters
+    db_session = AsyncMock()
+    aeros_db_session = AsyncMock()
+    project_name = "TestProject"
+    simulation_duration = 8760
+    overhaul_duration = 72
+    overhaul_interval = 8000
+    offset = 0
+    schematic_name = "TestSchematic"
+    
+    # Mock get_aeros_schematic_by_name
+    mock_schematic = MagicMock()
+    mock_schematic.id = 1
+    
+    # Mock db_session.execute
+    mock_eq = MagicMock()
+    mock_eq.location_tag = "EQ_MISSING"
+    mock_eq_result = MagicMock()
+    mock_eq_result.scalars.return_value.all.return_value = [mock_eq]
+    
+    mock_empty_result = MagicMock()
+    mock_empty_result.scalars.return_value.all.return_value = []
+
+    db_session.execute.side_effect = [
+        mock_eq_result,    # equipments
+        mock_empty_result, # rbd_group
+        mock_empty_result  # non_repairable
+    ]
+    
+    # Mock nodes data
+    mock_nodes_data = [
+        {"equipmentName": "EQ_MISSING", "relDisType": "Fixed", "relDisP1": 0}
+    ]
+    
+    with patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic), \
+         patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data), \
+         patch("src.aeros_equipment.service.get_asset_batch", return_value={}), \
+         patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None), \
+         patch("src.aeros_equipment.service.update_node", return_value=None) as mock_update_node:
+        
+        # Run
+        results = await update_equipment_for_simulation(
+            db_session=db_session,
+            aeros_db_session=aeros_db_session,
+            project_name=project_name,
+            simulation_duration=simulation_duration,
+            overhaul_duration=overhaul_duration,
+            overhaul_interval=overhaul_interval,
+            offset=offset,
+            schematic_name=schematic_name
+        )
+        
+        assert "EQ_MISSING" in results
+        # Default values when no reliability data
+        assert results["EQ_MISSING"]["mttr"] == 0
+        assert results["EQ_MISSING"]["distribution"] == "Fixed"
+        
+        # Check defaults in update_node call
+        called_args = mock_update_node.call_args[1]
+        eq_node = called_args["equipment_nodes"][0]
+        assert eq_node["cmDisUnitCode"] == "UHour"
+        assert eq_node["relDisUnitCode"] == "UHour"
+        assert eq_node["ohDisUnitCode"] == "UHour"