split up large test into smaller tests in test_evaluate.py

Author: ndawe

root_numpy/tests/test_evaluate.py

@@ -5,23 +5,17 @@
 from . import RNG
 
 
-def test_evaluate():
-    # create functions and histograms
+def test_evaluate_func():
     f1 = ROOT.TF1("f1", "x")
     f2 = ROOT.TF2("f2", "x*y")
     f3 = ROOT.TF3("f3", "x*y*z")
-    h1 = ROOT.TH1D("h1", "", 10, 0, 1)
-    h1.FillRandom("f1")
-    h2 = ROOT.TH2D("h2", "", 10, 0, 1, 10, 0, 1)
-    h2.FillRandom("f2")
-    h3 = ROOT.TH3D("h3", "", 10, 0, 1, 10, 0, 1, 10, 0, 1)
-    h3.FillRandom("f3")
+
     # generate random arrays
     arr_1d = RNG.rand(5)
     arr_2d = RNG.rand(5, 2)
     arr_3d = RNG.rand(5, 3)
     arr_4d = RNG.rand(5, 4)
-    # evaluate the functions
+
     assert_array_equal(rnp.evaluate(f1, arr_1d),
                        [f1.Eval(x) for x in arr_1d])
     assert_array_equal(rnp.evaluate(f1.GetTitle(), arr_1d),
@@ -38,14 +32,46 @@ def test_evaluate():
     f4 = ROOT.TFormula('test', 'x*y+z*t')
     assert_array_equal(rnp.evaluate(f4, arr_4d),
                        [f4.Eval(*x) for x in arr_4d])
-    # evaluate the histograms
+
+    assert_raises(ValueError, rnp.evaluate, f1, arr_2d)
+    assert_raises(ValueError, rnp.evaluate, f2, arr_3d)
+    assert_raises(ValueError, rnp.evaluate, f2, arr_1d)
+    assert_raises(ValueError, rnp.evaluate, f3, arr_1d)
+    assert_raises(ValueError, rnp.evaluate, f3, arr_2d)
+
+    assert_raises(ValueError, rnp.evaluate, "f", arr_1d)
+    assert_raises(ValueError, rnp.evaluate, "x*y", arr_1d)
+    assert_raises(ValueError, rnp.evaluate, "x", arr_2d)
+    assert_raises(ValueError, rnp.evaluate, "x*y", arr_3d)
+
+
+def test_evaluate_hist():
+    h1 = ROOT.TH1D("h1", "", 10, 0, 1)
+    h1.FillRandom("f1")
+    h2 = ROOT.TH2D("h2", "", 10, 0, 1, 10, 0, 1)
+    h2.FillRandom("f2")
+    h3 = ROOT.TH3D("h3", "", 10, 0, 1, 10, 0, 1, 10, 0, 1)
+    h3.FillRandom("f3")
+
+    arr_1d = RNG.rand(5)
+    arr_2d = RNG.rand(5, 2)
+    arr_3d = RNG.rand(5, 3)
+
     assert_array_equal(rnp.evaluate(h1, arr_1d),
                        [h1.GetBinContent(h1.FindBin(x)) for x in arr_1d])
     assert_array_equal(rnp.evaluate(h2, arr_2d),
                        [h2.GetBinContent(h2.FindBin(*x)) for x in arr_2d])
     assert_array_equal(rnp.evaluate(h3, arr_3d),
                        [h3.GetBinContent(h3.FindBin(*x)) for x in arr_3d])
-    # create a graph
+
+    assert_raises(ValueError, rnp.evaluate, h1, arr_2d)
+    assert_raises(ValueError, rnp.evaluate, h2, arr_3d)
+    assert_raises(ValueError, rnp.evaluate, h2, arr_1d)
+    assert_raises(ValueError, rnp.evaluate, h3, arr_1d)
+    assert_raises(ValueError, rnp.evaluate, h3, arr_2d)
+
+
+def test_evaluate_graph():
     g = ROOT.TGraph(2)
     g.SetPoint(0, 0, 1)
     g.SetPoint(1, 1, 2)
@@ -54,20 +80,7 @@ def test_evaluate():
     assert_array_equal(rnp.evaluate(s, [0, .5, 1]),
                        [s.Eval(x) for x in [0, .5, 1]])
     # test exceptions
+    arr_2d = RNG.rand(5, 2)
     assert_raises(TypeError, rnp.evaluate, object(), [1, 2, 3])
-    assert_raises(ValueError, rnp.evaluate, h1, arr_2d)
-    assert_raises(ValueError, rnp.evaluate, h2, arr_3d)
-    assert_raises(ValueError, rnp.evaluate, h2, arr_1d)
-    assert_raises(ValueError, rnp.evaluate, h3, arr_1d)
-    assert_raises(ValueError, rnp.evaluate, h3, arr_2d)
-    assert_raises(ValueError, rnp.evaluate, f1, arr_2d)
-    assert_raises(ValueError, rnp.evaluate, f2, arr_3d)
-    assert_raises(ValueError, rnp.evaluate, f2, arr_1d)
-    assert_raises(ValueError, rnp.evaluate, f3, arr_1d)
-    assert_raises(ValueError, rnp.evaluate, f3, arr_2d)
     assert_raises(ValueError, rnp.evaluate, g, arr_2d)
     assert_raises(ValueError, rnp.evaluate, s, arr_2d)
-    assert_raises(ValueError, rnp.evaluate, "f", arr_1d)
-    assert_raises(ValueError, rnp.evaluate, "x*y", arr_1d)
-    assert_raises(ValueError, rnp.evaluate, "x", arr_2d)
-    assert_raises(ValueError, rnp.evaluate, "x*y", arr_3d)