From 3bc9f2286b039b7bcf8daf03bec6d5688d444dcd Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 20 Feb 2026 20:28:12 +0100
Subject: [PATCH] fix(test): update dqn_training_pipeline_test callback to
 3-arg signature

The train() callback was changed to (epoch, data, is_best) but
this test still used (epoch, data). Update all 5 call sites.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 ml/tests/dqn_training_pipeline_test.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/ml/tests/dqn_training_pipeline_test.rs b/ml/tests/dqn_training_pipeline_test.rs
index 6373c7e3e..3e180241d 100644
--- a/ml/tests/dqn_training_pipeline_test.rs
+++ b/ml/tests/dqn_training_pipeline_test.rs
@@ -105,7 +105,7 @@ async fn test_dqn_trains_on_es_fut() -> Result<()> {
     let mut final_checkpoint_path = PathBuf::new();
 
     let metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             let path = checkpoint_dir.join(format!("dqn_test_epoch_{}.safetensors", epoch));
             std::fs::write(&path, checkpoint_data)?;
             checkpoint_saved = true;
@@ -224,7 +224,7 @@ async fn test_dqn_loss_decreases() -> Result<()> {
 
     // Track losses per epoch (would need to modify trainer to expose this)
     let metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             let path = checkpoint_dir.join(format!("dqn_loss_test_epoch_{}.safetensors", epoch));
             std::fs::write(&path, checkpoint_data)?;
             Ok(path.to_string_lossy().to_string())
@@ -282,7 +282,7 @@ async fn test_dqn_checkpoint_save_load() -> Result<()> {
     let mut saved_checkpoint_path = PathBuf::new();
 
     let _metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             let path =
                 checkpoint_dir.join(format!("dqn_checkpoint_test_epoch_{}.safetensors", epoch));
             std::fs::write(&path, checkpoint_data)?;
@@ -345,7 +345,7 @@ async fn test_dqn_q_value_predictions() -> Result<()> {
     let mut trainer = DQNTrainer::new(hyperparams)?;
 
     let metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             let path = checkpoint_dir.join(format!("dqn_qvalue_test_epoch_{}.safetensors", epoch));
             std::fs::write(&path, checkpoint_data)?;
             Ok(path.to_string_lossy().to_string())
@@ -401,7 +401,7 @@ async fn test_dqn_epsilon_greedy() -> Result<()> {
     let mut trainer = DQNTrainer::new(hyperparams)?;
 
     let metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             let path = checkpoint_dir.join(format!("dqn_epsilon_test_epoch_{}.safetensors", epoch));
             std::fs::write(&path, checkpoint_data)?;
             Ok(path.to_string_lossy().to_string())
@@ -470,7 +470,7 @@ async fn test_dqn_full_production_training() -> Result<()> {
     let mut epoch_count = 0;
 
     let metrics = trainer
-        .train(&data_dir, |epoch, checkpoint_data| {
+        .train(&data_dir, |epoch, checkpoint_data, _is_best| {
             epoch_count += 1;
             let path = if epoch == hyperparams.epochs {
                 production_checkpoint_path.clone()