From f2c51636aa6171279944df9cd837471171c8f38e Mon Sep 17 00:00:00 2001 From: Cizz22 Date: Sun, 19 Apr 2026 12:28:42 +0700 Subject: [PATCH] refactor: return equipment nodes and reliability data from update_equipment_for_simulation instead of performing side-effect updates --- src/aeros_equipment/service.py | 11 ++--- src/aeros_simulation/router.py | 6 ++- temporal/workflow.py | 4 +- .../unit/test_update_equipment_simulation.py | 41 ++++++++++--------- 4 files changed, 33 insertions(+), 29 deletions(-) diff --git a/src/aeros_equipment/service.py b/src/aeros_equipment/service.py index 315746b..0c8145d 100644 --- a/src/aeros_equipment/service.py +++ b/src/aeros_equipment/service.py @@ -622,13 +622,10 @@ async def update_equipment_for_simulation( ) reqNodeInputs.append(eq) - # await update_node( - # db_session=db_session, - # equipment_nodes=reqNodeInputs, - # project_name=project_name - # ) - - return reqNodeInputs + return { + "equipment_nodes": reqNodeInputs, + "reliability": results + } # Optimized individual fetch functions async def get_equipment_mttr(*, location_tag: str, client: httpx.AsyncClient) -> float: diff --git a/src/aeros_simulation/router.py b/src/aeros_simulation/router.py index dbfbb76..d027849 100644 --- a/src/aeros_simulation/router.py +++ b/src/aeros_simulation/router.py @@ -717,7 +717,11 @@ async def calculate_contribution( # ) await execute_simulation( - db_session=db_session, simulation_id=simulation.id, sim_data=sim_data, is_saved=True, eq_update=results + db_session=db_session, + simulation_id=simulation.id, + sim_data=sim_data, + is_saved=True, + eq_update=results['reliability'] ) eaf, edh, efficiency_uptime = await calculate_plant_eaf(db_session=db_session, simulation_id=simulation.id) diff --git a/temporal/workflow.py b/temporal/workflow.py index 3caafd7..3f5b2c7 100644 --- a/temporal/workflow.py +++ b/temporal/workflow.py @@ -41,14 +41,14 @@ class SimulationWorkflow: await workflow.execute_activity( execute_update_node, - {"projectName": sim_data['projectName'], "equipment_nodes": results}, + {"projectName": sim_data['projectName'], "equipment_nodes": results['equipment_nodes']}, start_to_close_timeout=timedelta(days=1) ) # 2. Execute simulation await workflow.execute_activity( execute_simulation_activity, - {"sim_data": sim_data, "eq_update": results}, + {"sim_data": sim_data, "eq_update": results['reliability']}, start_to_close_timeout=timedelta(days=1) ) diff --git a/tests/unit/test_update_equipment_simulation.py b/tests/unit/test_update_equipment_simulation.py index b5b511b..f2a6d75 100644 --- a/tests/unit/test_update_equipment_simulation.py +++ b/tests/unit/test_update_equipment_simulation.py @@ -63,10 +63,10 @@ async def test_update_equipment_for_simulation_basic(): 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: + patch("src.aeros_equipment.service.update_node", return_value=None): # Run the function - results = await update_equipment_for_simulation( + res = await update_equipment_for_simulation( db_session=db_session, aeros_db_session=aeros_db_session, project_name=project_name, @@ -77,15 +77,14 @@ async def test_update_equipment_for_simulation_basic(): schematic_name=schematic_name ) + results = res['reliability'] + equipment_nodes = res['equipment_nodes'] + # 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" @@ -144,18 +143,20 @@ async def test_update_equipment_for_simulation_custom_input(): patch("src.aeros_equipment.service.update_node", return_value=None): # Run - results = await update_equipment_for_simulation( + res = 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, + overhaul_interval=8000, + offset=0, schematic_name=schematic_name, custom_input=custom_input ) + results = res['reliability'] + assert "EQ001" in results assert results["EQ001"]["mttr"] == 10 assert results["EQ001"]["distribution"] == "Exponential2" @@ -200,10 +201,10 @@ async def test_update_equipment_for_simulation_missing_reliability(): 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: + patch("src.aeros_equipment.service.update_node", return_value=None): # Run - results = await update_equipment_for_simulation( + res = await update_equipment_for_simulation( db_session=db_session, aeros_db_session=aeros_db_session, project_name=project_name, @@ -214,14 +215,16 @@ async def test_update_equipment_for_simulation_missing_reliability(): schematic_name=schematic_name ) + results = res['reliability'] + equipment_nodes = res['equipment_nodes'] + 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] + # Check defaults in returned equipment_nodes + eq_node = equipment_nodes[0] assert eq_node["cmDisUnitCode"] == "UHour" assert eq_node["relDisUnitCode"] == "UHour" assert eq_node["ohDisUnitCode"] == "UHour" @@ -278,10 +281,10 @@ async def test_update_equipment_for_simulation_sigma_constraint(): 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: + patch("src.aeros_equipment.service.update_node", return_value=None): # Run - await update_equipment_for_simulation( + res = await update_equipment_for_simulation( db_session=db_session, aeros_db_session=aeros_db_session, project_name=project_name, @@ -292,7 +295,7 @@ async def test_update_equipment_for_simulation_sigma_constraint(): schematic_name=schematic_name ) - # Check if sigma was constrained correctly - called_args = mock_update_node.call_args[1] - eq_node = called_args["equipment_nodes"][0] + # Check if sigma was constrained correctly in returned equipment_nodes + equipment_nodes = res['equipment_nodes'] + eq_node = equipment_nodes[0] assert eq_node["relDisP2"] == 999