use pytest.mark.parametrize

Author: basnijholt

adaptive/tests/test_balancing_learner.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+import pytest
+
 from adaptive.learner import Learner1D, BalancingLearner
 from adaptive.runner import simple
 
@@ -24,32 +26,30 @@ def test_balancing_learner_loss_cache():
     assert bl.loss(real=True) == real_loss
 
 
-def test_distribute_first_points_over_learners():
-    for strategy in ['loss', 'loss_improvements', 'npoints']:
-        learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
-        learner = BalancingLearner(learners, strategy=strategy)
-        points, _ = learner.ask(100)
-        i_learner, xs = zip(*points)
-        # assert that are all learners in the suggested points
-        assert len(set(i_learner)) == len(learners)
-
-
-def test_ask_0():
-    for strategy in ['loss', 'loss_improvements', 'npoints']:
-        learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
-        learner = BalancingLearner(learners, strategy=strategy)
-        points, _ = learner.ask(0)
-        assert len(points) == 0
-
-
-def test_strategies():
-    goals = {
-        'loss': lambda l: l.loss() < 0.1,
-        'loss_improvements': lambda l: l.loss() < 0.1,
-        'npoints': lambda bl: all(l.npoints > 10 for l in bl.learners)
-    }
-
-    for strategy, goal in goals.items():
-        learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
-        learner = BalancingLearner(learners, strategy=strategy)
-        simple(learner, goal=goal)
+@pytest.mark.parametrize('strategy', ['loss', 'loss_improvements', 'npoints'])
+def test_distribute_first_points_over_learners(strategy):
+    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+    learner = BalancingLearner(learners, strategy=strategy)
+    points, _ = learner.ask(100)
+    i_learner, xs = zip(*points)
+    # assert that are all learners in the suggested points
+    assert len(set(i_learner)) == len(learners)
+
+
+@pytest.mark.parametrize('strategy', ['loss', 'loss_improvements', 'npoints'])
+def test_ask_0(strategy):
+    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+    learner = BalancingLearner(learners, strategy=strategy)
+    points, _ = learner.ask(0)
+    assert len(points) == 0
+
+
+@pytest.mark.parametrize('strategy, goal', [
+    ('loss', lambda l: l.loss() < 0.1),
+    ('loss_improvements', lambda l: l.loss() < 0.1),
+    ('npoints', lambda bl: all(l.npoints > 10 for l in bl.learners)),
+])
+def test_strategies(strategy, goal):
+    learners = [Learner1D(lambda x: x, bounds=(-1, 1)) for i in range(10)]
+    learner = BalancingLearner(learners, strategy=strategy)
+    simple(learner, goal=goal)