From 7bae73c60da37932851f8373cc567f8f35e8a078 Mon Sep 17 00:00:00 2001
From: Brady Planden <brady.planden@gmail.com>
Date: Sat, 20 Jul 2024 14:13:38 +0100
Subject: [PATCH] tests: increase coverage, fix flaky threshold test.

---
 .gitignore                      |  3 +++
 tests/unit/test_optimisation.py | 17 +++++++++++++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3c3bb708d..2bafcca63 100644
--- a/.gitignore
+++ b/.gitignore
@@ -314,3 +314,6 @@ $RECYCLE.BIN/
 # Airspeed Velocity
 *.asv/
 results/
+
+# Pycharm
+*.idea/
diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py
index bb6d08fa2..6d5a035cf 100644
--- a/tests/unit/test_optimisation.py
+++ b/tests/unit/test_optimisation.py
@@ -1,3 +1,5 @@
+import io
+import sys
 import warnings
 
 import numpy as np
@@ -422,12 +424,23 @@ def test_halting(self, cost):
         with pytest.raises(ValueError):
             optim.set_max_evaluations(-1)
 
-        optim = pybop.Optimisation(cost=cost)
+        # Reset optim
+        optim = pybop.Optimisation(cost=cost, sigma0=0.015, verbose=True)
 
-        # Trigger threshold
+        # Confirm setting threshold == None
         optim.set_threshold(None)
+        assert optim._threshold is None
+
+        # Confirm threshold halts
+        # Redirect stdout to capture print output
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
         optim.set_threshold(np.inf)
         optim.run()
+        assert (
+            captured_output.getvalue().strip()
+            == "Halt: Objective function crossed threshold: inf."
+        )
         optim.set_max_unchanged_iterations()
 
         # Test callback and halting output