refactor: remove shaft from process system

Author: jsolaas

src/libecalc/domain/process/process_solver/solvers/speed_solver.py

@@ -4,6 +4,7 @@
     find_root,
     maximize_x_given_boolean_condition_function,
 )
+from libecalc.domain.process.entities.shaft import Shaft
 from libecalc.domain.process.process_solver.boundary import Boundary
 from libecalc.domain.process.process_solver.solver import Solver
 from libecalc.domain.process.process_system.process_error import ProcessError
@@ -14,16 +15,17 @@
 
 
 class SpeedSolver(Solver):
-    def __init__(self, boundary: Boundary, target_pressure: float):
+    def __init__(self, boundary: Boundary, target_pressure: float, shaft: Shaft):
         self._boundary = boundary
         self._target_pressure = target_pressure
+        self._shaft = shaft
 
     def solve(
         self,
         process_system: ProcessSystem,
         inlet_stream: FluidStream,
     ) -> FluidStream | None:
-        shaft = process_system.get_shaft()
+        shaft = self._shaft
 
         def get_outlet_stream(speed: float) -> FluidStream:
             shaft.set_speed(speed)

src/libecalc/domain/process/process_system/process_system.py

@@ -9,12 +9,10 @@
 class ProcessSystem:
     def __init__(
         self,
-        shaft: Shaft,
         process_units: list[ProcessUnit],
         upstream_choke: Choke | None = None,
         downstream_choke: Choke | None = None,
     ):
-        self._shaft = shaft
         self._process_units = process_units
         self._upstream_choke = upstream_choke
         self._downstream_choke = downstream_choke

tests/libecalc/domain/process/conftest.py

@@ -2,7 +2,7 @@
 
 from libecalc.domain.process.compressor.core.train.stage import CompressorTrainStage
 from libecalc.domain.process.entities.process_units.choke import Choke
-from libecalc.domain.process.entities.shaft import Shaft, VariableSpeedShaft
+from libecalc.domain.process.entities.shaft import Shaft
 from libecalc.domain.process.process_solver.stream_constraint import PressureStreamConstraint
 from libecalc.domain.process.process_system.process_error import OutsideCapacityError
 from libecalc.domain.process.process_system.process_system import ProcessSystem
@@ -191,13 +191,10 @@ def create_process_system(
         process_units: list[ProcessUnit] | None = None,
         downstream_choke: Choke | None = None,
         upstream_choke: Choke | None = None,
-        shaft: Shaft | None = None,
     ):
         if process_units is None:
             process_units = [simple_process_unit_factory(fluid_service)]
-        shaft = shaft or VariableSpeedShaft()
         return ProcessSystem(
-            shaft=shaft,
             process_units=process_units,
             downstream_choke=downstream_choke,
             upstream_choke=upstream_choke,

tests/libecalc/domain/process/solvers/test_speed_solver.py

@@ -49,15 +49,14 @@ def test_speed_solver(
     speed_solver = SpeedSolver(
         boundary=speed_boundary,
         target_pressure=target_pressure,
+        shaft=shaft,
     )
     inlet_stream = stream_factory(
         standard_rate_m3_per_day=1000,
         pressure_bara=inlet_pressure,
     )
 
-    process_system = process_system_factory(
-        shaft=shaft, process_units=[SpeedProcessUnit(shaft=shaft, fluid_service=fluid_service)]
-    )
+    process_system = process_system_factory(process_units=[SpeedProcessUnit(shaft=shaft, fluid_service=fluid_service)])
 
     outlet_stream = speed_solver.solve(process_system=process_system, inlet_stream=inlet_stream)
 

tests/libecalc/domain/process/test_process_solver.py

@@ -37,7 +37,6 @@ def test_speed_solver(
     target_pressure = 90
     stream_constraint = stream_constraint_factory(pressure=target_pressure)
     process_system = process_system_factory(
-        shaft=shaft,
         process_units=[compressor_train_stage_process_unit_factory(chart_data=chart_data, shaft=shaft)],
         downstream_choke=Choke(fluid_service=fluid_service),
     )
@@ -49,12 +48,13 @@ def test_speed_solver(
             SpeedSolver(
                 target_pressure=target_pressure,
                 boundary=Boundary(min=min(chart_speeds), max=max(chart_speeds)),
+                shaft=shaft,
             ),
         ],
         stream_constraint=stream_constraint,
     )
     assert process_solver.find_solution()
-    assert process_system.get_shaft().get_speed() == snapshot(102.66555959304989)
+    assert shaft.get_speed() == snapshot(102.66555959304989)
 
     outlet_stream = process_system.propagate_stream(inlet_stream=inlet_stream)
     assert outlet_stream.pressure_bara == pytest.approx(target_pressure)