From 57593f6a28e914a7de648595be4487b6b28955c3 Mon Sep 17 00:00:00 2001 From: Cizz22 Date: Sun, 19 Apr 2026 11:36:02 +0700 Subject: [PATCH] refactor: move node update logic to a dedicated Temporal activity and update workflow execution sequence --- src/aeros_equipment/service.py | 12 +- temporal/activity.py | 10 + temporal/workflow.py | 8 +- .../unit/test_update_equipment_simulation.py | 202 +++++++++++------- 4 files changed, 152 insertions(+), 80 deletions(-) diff --git a/src/aeros_equipment/service.py b/src/aeros_equipment/service.py index 291f120..64d9a80 100644 --- a/src/aeros_equipment/service.py +++ b/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: diff --git a/temporal/activity.py b/temporal/activity.py index 842f26b..d692901 100644 --- a/temporal/activity.py +++ b/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): diff --git a/temporal/workflow.py b/temporal/workflow.py index bcfab7d..3caafd7 100644 --- a/temporal/workflow.py +++ b/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, diff --git a/tests/unit/test_update_equipment_simulation.py b/tests/unit/test_update_equipment_simulation.py index 253e0fd..c6fa078 100644 --- a/tests/unit/test_update_equipment_simulation.py +++ b/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,44 +58,43 @@ 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) - - # 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" + + 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(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,31 +132,96 @@ 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) - - # 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" + + 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"