add type annotations for adaptive/tests/test_balancing_learner.py

Author: basnijholt

adaptive/tests/test_balancing_learner.py

@@ -1,3 +1,5 @@
+from typing import Callable
+
 import pytest
 
 from adaptive.learner import BalancingLearner, Learner1D
@@ -6,7 +8,7 @@
 strategies = ["loss", "loss_improvements", "npoints", "cycle"]
 
 
-def test_balancing_learner_loss_cache():
+def test_balancing_learner_loss_cache() -> None:
     learner = Learner1D(lambda x: x, bounds=(-1, 1))
     learner.tell(-1, -1)
     learner.tell(1, 1)
@@ -27,7 +29,7 @@ def test_balancing_learner_loss_cache():
 
 
 @pytest.mark.parametrize("strategy", strategies)
-def test_distribute_first_points_over_learners(strategy):
+def test_distribute_first_points_over_learners(strategy: str) -> None:
     for initial_points in [0, 3]:
         learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
         learner = BalancingLearner(learners, strategy=strategy)
@@ -42,7 +44,7 @@ def test_distribute_first_points_over_learners(strategy):
 
 
 @pytest.mark.parametrize("strategy", strategies)
-def test_ask_0(strategy):
+def test_ask_0(strategy: str) -> None:
     learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
     learner = BalancingLearner(learners, strategy=strategy)
     points, _ = learner.ask(0)
@@ -58,7 +60,7 @@ def test_ask_0(strategy):
         ("cycle", lambda l: l.loss() < 0.1),
     ],
 )
-def test_strategies(strategy, goal):
+def test_strategies(strategy: str, goal: Callable) -> None:
     learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
     learner = BalancingLearner(learners, strategy=strategy)
     simple(learner, goal=goal)