editoast: move get_towed_rolling_stock to STDCMRequestPayload

Signed-off-by: hamz2a <atrari.hamza@gmail.com>
OpenRailAssociation · Nov 4, 2024 · 0e3603c · 0e3603c
1 parent cf8d096
commit 0e3603c
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 16 deletions.
diff --git a/editoast/src/views/timetable/stdcm.rs b/editoast/src/views/timetable/stdcm.rs
@@ -36,7 +36,6 @@ use crate::core::AsCoreRequest;
 use crate::core::CoreClient;
 use crate::error::Result;
 use crate::models::timetable::TimetableWithTrains;
-use crate::models::towed_rolling_stock::TowedRollingStockModel;
 use crate::models::train_schedule::TrainSchedule;
 use crate::models::Infra;
 use crate::models::RollingStockModel;
@@ -158,20 +157,6 @@ async fn stdcm(
         })
         .await?;
 
-    let towed_rolling_stock =
-        if let Some(towed_rolling_stock_id) = stdcm_request.towed_rolling_stock_id {
-            let towed_rolling_stock =
-                TowedRollingStockModel::retrieve_or_fail(conn, towed_rolling_stock_id, || {
-                    STDCMError::TowedRollingStockNotFound {
-                        towed_rolling_stock_id,
-                    }
-                })
-                .await?;
-            Some(towed_rolling_stock)
-        } else {
-            None
-        };
-
     // 2. Compute the earliest start time and maximum departure delay
     let (virtual_train_schedule, virtual_train_sim_result, virtual_train_pathfinding_result) =
         simulate_train_run(
@@ -254,7 +239,10 @@ async fn stdcm(
             max_speed: stdcm_request.max_speed,
             total_length: stdcm_request.total_length,
             total_mass: stdcm_request.total_mass,
-            towed_rolling_stock: towed_rolling_stock.map(From::from),
+            towed_rolling_stock: stdcm_request
+                .get_towed_rolling_stock(conn)
+                .await?
+                .map(From::from),
             traction_engine: rolling_stock.into(),
         }
         .into(),

diff --git a/editoast/src/views/timetable/stdcm_request_payload.rs b/editoast/src/views/timetable/stdcm_request_payload.rs
@@ -17,11 +17,13 @@ use crate::core::stdcm::STDCMPathItem;
 use crate::core::stdcm::STDCMStepTimingData;
 use crate::error::Result;
 use crate::models::temporary_speed_limits::TemporarySpeedLimit;
+use crate::models::towed_rolling_stock::TowedRollingStockModel;
 use crate::models::work_schedules::WorkSchedule;
 use crate::models::List;
 use crate::views::path::path_item_cache::PathItemCache;
 use crate::views::path::pathfinding::PathfindingFailure;
 use crate::views::path::pathfinding::PathfindingResult;
+use crate::Retrieve;
 use crate::SelectionSettings;
 
 use super::stdcm::STDCMError;
@@ -261,4 +263,23 @@ impl STDCMRequestPayload {
             .filter(move || WorkSchedule::WORK_SCHEDULE_GROUP_ID.eq(work_schedule_group_id));
         WorkSchedule::list(conn, selection_setting).await
     }
+
+    pub async fn get_towed_rolling_stock(
+        &self,
+        conn: &mut DbConnection,
+    ) -> Result<Option<TowedRollingStockModel>> {
+        if self.towed_rolling_stock_id.is_none() {
+            return Ok(None);
+        }
+
+        let towed_rolling_stock_id = self.towed_rolling_stock_id.unwrap();
+        let towed_rolling_stock =
+            TowedRollingStockModel::retrieve_or_fail(conn, towed_rolling_stock_id, || {
+                STDCMError::TowedRollingStockNotFound {
+                    towed_rolling_stock_id,
+                }
+            })
+            .await?;
+        Ok(Some(towed_rolling_stock))
+    }
 }