diff --git a/crates/ml/src/trainers/dqn/trainer/tests.rs b/crates/ml/src/trainers/dqn/trainer/tests.rs
index 363361a7e..8bb396d5c 100644
--- a/crates/ml/src/trainers/dqn/trainer/tests.rs
+++ b/crates/ml/src/trainers/dqn/trainer/tests.rs
@@ -1,4 +1,5 @@
 use super::*;
+use crate::dqn::TradingState;
 use crate::hyperopt::ParameterSpace;
 use std::sync::OnceLock;