From 7d0fc228099113107eacaadda83f42cc91c319b8 Mon Sep 17 00:00:00 2001
From: Bas Nijholt <basnijholt@gmail.com>
Date: Wed, 18 Sep 2019 17:09:38 +0200
Subject: [PATCH] run test_balancing_learner for all strategies

---
 adaptive/tests/test_balancing_learner.py | 57 +++++++++++++++++++++++-
 adaptive/tests/test_learners.py          | 46 +------------------
 2 files changed, 57 insertions(+), 46 deletions(-)

diff --git a/adaptive/tests/test_balancing_learner.py b/adaptive/tests/test_balancing_learner.py
index fff0ff186..c124dbce8 100644
--- a/adaptive/tests/test_balancing_learner.py
+++ b/adaptive/tests/test_balancing_learner.py
@@ -1,10 +1,20 @@
 # -*- coding: utf-8 -*-
 
+import random
+
 import pytest
 
-from adaptive.learner import BalancingLearner, Learner1D
+from adaptive.learner import (
+    BalancingLearner,
+    Learner1D,
+    Learner2D,
+    LearnerND,
+    AverageLearner,
+    SequenceLearner,
+)
 from adaptive.runner import simple
 
+from .test_learners import generate_random_parametrization, run_with
 
 strategies = ["loss", "loss_improvements", "npoints", "cycle"]
 
@@ -65,3 +75,48 @@ 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)
+
+
+@run_with(
+    Learner1D,
+    Learner2D,
+    LearnerND,
+    AverageLearner,
+    SequenceLearner,
+    with_all_loss_functions=False,
+)
+@pytest.mark.parametrize("strategy", ["loss", "loss_improvements", "npoints", "cycle"])
+def test_balancing_learner(learner_type, f, learner_kwargs, strategy):
+    """Test if the BalancingLearner works with the different types of learners."""
+    learners = [
+        learner_type(generate_random_parametrization(f), **learner_kwargs)
+        for i in range(4)
+    ]
+
+    learner = BalancingLearner(learners, strategy=strategy)
+
+    # Emulate parallel execution
+    stash = []
+
+    for i in range(100):
+        n = random.randint(1, 10)
+        m = random.randint(0, n)
+        xs, _ = learner.ask(n, tell_pending=False)
+
+        # Save 'm' random points out of `xs` for later
+        random.shuffle(xs)
+        for _ in range(m):
+            stash.append(xs.pop())
+
+        for x in xs:
+            learner.tell(x, learner.function(x))
+
+        # Evaluate and add 'm' random points from `stash`
+        random.shuffle(stash)
+        for _ in range(m):
+            x = stash.pop()
+            learner.tell(x, learner.function(x))
+
+    assert all(l.npoints > 10 for l in learner.learners), [
+        l.npoints for l in learner.learners
+    ]
diff --git a/adaptive/tests/test_learners.py b/adaptive/tests/test_learners.py
index bf84fcc14..371efc61e 100644
--- a/adaptive/tests/test_learners.py
+++ b/adaptive/tests/test_learners.py
@@ -30,7 +30,7 @@
 
 try:
     from adaptive.learner import SKOptLearner
-except ModuleNotFoundError:
+except (ModuleNotFoundError, ImportError):
     SKOptLearner = None
 
 
@@ -456,50 +456,6 @@ def test_learner_performance_is_invariant_under_scaling(
     assert math.isclose(learner.loss(), control.loss(), rel_tol=1e-10)
 
 
-@run_with(
-    Learner1D,
-    Learner2D,
-    LearnerND,
-    AverageLearner,
-    SequenceLearner,
-    with_all_loss_functions=False,
-)
-def test_balancing_learner(learner_type, f, learner_kwargs):
-    """Test if the BalancingLearner works with the different types of learners."""
-    learners = [
-        learner_type(generate_random_parametrization(f), **learner_kwargs)
-        for i in range(4)
-    ]
-
-    learner = BalancingLearner(learners)
-
-    # Emulate parallel execution
-    stash = []
-
-    for i in range(100):
-        n = random.randint(1, 10)
-        m = random.randint(0, n)
-        xs, _ = learner.ask(n, tell_pending=False)
-
-        # Save 'm' random points out of `xs` for later
-        random.shuffle(xs)
-        for _ in range(m):
-            stash.append(xs.pop())
-
-        for x in xs:
-            learner.tell(x, learner.function(x))
-
-        # Evaluate and add 'm' random points from `stash`
-        random.shuffle(stash)
-        for _ in range(m):
-            x = stash.pop()
-            learner.tell(x, learner.function(x))
-
-    assert all(l.npoints > 10 for l in learner.learners), [
-        l.npoints for l in learner.learners
-    ]
-
-
 @run_with(
     Learner1D,
     Learner2D,