Replace assertLogs

Author: albertvillanova

tests/test_grpo_trainer.py

@@ -1027,7 +1027,7 @@ def test_training_with_mask_truncated_completions_all_masked(self):
             new_param = trainer.model.get_parameter(n)
             assert torch.equal(param, new_param), f"Parameter {n} has changed."
 
-    def test_warning_raised_all_rewards_none(self):
+    def test_warning_raised_all_rewards_none(self, caplog):
         """Test that a proper warning is raised when all rewards are None."""
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
 
@@ -1050,11 +1050,11 @@ def always_none_reward_func(completions, **kwargs):
             train_dataset=dataset,
         )
 
-        with self.assertLogs("trl.trainer.grpo_trainer", level="WARNING") as cm:
+        with caplog.at_level("WARNING", logger="trl.trainer.grpo_trainer"):
             trainer.train()
 
         expected_warning = "All reward functions returned None for the following kwargs:"
-        assert expected_warning in cm.output[0]
+        assert expected_warning in caplog.text
 
     def test_training_num_generations_larger_than_batch_size(self):
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")

tests/test_rloo_trainer.py

@@ -882,7 +882,7 @@ def test_training_with_mask_truncated_completions_all_masked(self):
             new_param = trainer.model.get_parameter(n)
             assert torch.equal(param, new_param), f"Parameter {n} has changed."
 
-    def test_warning_raised_all_rewards_none(self):
+    def test_warning_raised_all_rewards_none(self, caplog):
         """Test that a proper warning is raised when all rewards are None."""
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
 
@@ -905,11 +905,11 @@ def always_none_reward_func(completions, **kwargs):
             train_dataset=dataset,
         )
 
-        with self.assertLogs("trl.trainer.rloo_trainer", level="WARNING") as cm:
+        with caplog.at_level("WARNING", logger="trl.trainer.rloo_trainer"):
             trainer.train()
 
         expected_warning = "All reward functions returned None for the following kwargs:"
-        assert expected_warning in cm.output[0]
+        assert expected_warning in caplog.text
 
     def test_training_num_generations_larger_than_batch_size(self):
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")

tests/test_utils.py

@@ -422,7 +422,7 @@ def test_token_classification_task_with_ignored_tokens_1(self):
         result = compute_accuracy(eval_pred)
         assert round(abs(result["accuracy"] - expected_accuracy), 7) == 0
 
-    def test_rewards_comparison_task(self):
+    def test_rewards_comparison_task(self, caplog):
         eval_pred = (
             np.array(
                 [
@@ -435,15 +435,15 @@ def test_rewards_comparison_task(self):
         )
         expected_accuracy = 0.5  # 1 match, 1 mismatch, 1 equal (ignored)
 
-        with self.assertLogs("trl.trainer.utils", level="WARNING") as cm:
+        with caplog.at_level("WARNING", logger="trl.trainer.utils"):
             result = compute_accuracy(eval_pred)
 
         assert round(abs(result["accuracy"] - expected_accuracy), 7) == 0
         expected_warning = (
             "There are 1 out of 3 instances where the predictions for both options are equal. "
             "These instances are ignored in the accuracy computation."
         )
-        assert expected_warning in cm.output[0]
+        assert expected_warning in caplog.text
 
 
 class TestFlushLeft(TrlTestCase):