refactor: move node update logic to a dedicated Temporal activity and update workflow execution sequence

src/aeros_equipment/service.py

@@ -618,13 +618,13 @@ async def update_equipment_for_simulation(
             )
             reqNodeInputs.append(eq)
 
-    await update_node(
-        db_session=db_session,
-        equipment_nodes=reqNodeInputs,
-        project_name=project_name
-    )
+    # await update_node(
+    #     db_session=db_session,
+    #     equipment_nodes=reqNodeInputs,
+    #     project_name=project_name
+    # )
 
-    return results
+    return reqNodeInputs    
 
 # Optimized individual fetch functions
 async def get_equipment_mttr(*, location_tag: str, client: httpx.AsyncClient) -> float:

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

@@ -6,7 +6,7 @@ from src.aeros_simulation.model import AerosNode
 import logging
 
 @pytest.mark.asyncio
-async def test_update_equipment_for_simulation_basic(mocker):
+async def test_update_equipment_for_simulation_basic():
     # Mock parameters
     db_session = AsyncMock()
     aeros_db_session = AsyncMock()
@@ -20,23 +20,18 @@ async def test_update_equipment_for_simulation_basic(mocker):
     # Mock get_aeros_schematic_by_name
     mock_schematic = MagicMock()
     mock_schematic.id = 1
-    mocker.patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic)
     
-    # Mock db_session.execute for equipments
+    # 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 db_session.execute for rbd_group
     mock_group = MagicMock()
     mock_group.group_name = "GROUP001"
-    
     mock_group_result = MagicMock()
     mock_group_result.scalars.return_value.all.return_value = [mock_group]
     
-    # Mock db_session.execute for non_repairable
     mock_nr_result = MagicMock()
     mock_nr_result.scalars.return_value.all.return_value = []
 
@@ -46,14 +41,13 @@ async def test_update_equipment_for_simulation_basic(mocker):
         mock_nr_result     # for non_repairable
     ]
     
-    # Mock get_aeros_equipment_by_location_tag
+    # Mock nodes data
     mock_nodes_data = [
         {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0},
         {"equipmentName": "GROUP001", "relDisType": "Fixed", "relDisP1": 0}
     ]
-    mocker.patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data)
     
-    # Mock get_asset_batch
+    # Mock reliability data
     mock_reliability_data = {
         "EQ001": {
             "relDisType": "Weibull2",
@@ -64,13 +58,12 @@ async def test_update_equipment_for_simulation_basic(mocker):
             "cmDisP2": 0
         }
     }
-    mocker.patch("src.aeros_equipment.service.get_asset_batch", return_value=mock_reliability_data)
 
-    # Mock update_oh_interval_offset
-    mocker.patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None)
-    
-    # Mock update_node
-    mock_update_node = mocker.patch("src.aeros_equipment.service.update_node", return_value=None)
+    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(
@@ -101,7 +94,7 @@ async def test_update_equipment_for_simulation_basic(mocker):
         assert eq001_node["ohDisUnitCode"] == "UHour"
 
 @pytest.mark.asyncio
-async def test_update_equipment_for_simulation_custom_input(mocker):
+async def test_update_equipment_for_simulation_custom_input():
     # Mock parameters
     db_session = AsyncMock()
     aeros_db_session = AsyncMock()
@@ -123,12 +116,10 @@ async def test_update_equipment_for_simulation_custom_input(mocker):
     # Mock get_aeros_schematic_by_name
     mock_schematic = MagicMock()
     mock_schematic.id = 1
-    mocker.patch("src.aeros_equipment.service.get_aeros_schematic_by_name", return_value=mock_schematic)
     
     # 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]
     
@@ -141,17 +132,16 @@ async def test_update_equipment_for_simulation_custom_input(mocker):
         mock_empty_result  # non_repairable
     ]
     
-    # Mock get_aeros_equipment_by_location_tag
+    # Mock nodes data
     mock_nodes_data = [
         {"equipmentName": "EQ001", "relDisType": "Fixed", "relDisP1": 0}
     ]
-    mocker.patch("src.aeros_equipment.service.get_aeros_equipment_by_location_tag", return_value=mock_nodes_data)
     
-    # Mock get_asset_batch (should be called but maybe ignored if custom input matches)
-    mocker.patch("src.aeros_equipment.service.get_asset_batch", return_value={})
-    
-    mocker.patch("src.aeros_equipment.service.update_oh_interval_offset", return_value=None)
-    mock_update_node = mocker.patch("src.aeros_equipment.service.update_node", return_value=None)
+    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(
@@ -169,3 +159,69 @@ async def test_update_equipment_for_simulation_custom_input(mocker):
         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"