added separate test for nanmean_f, reset the order of pytest fixtures

neurodata · Aug 21, 2024 · 251392f · 251392f
1 parent f852fc4
commit 251392f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/treeple/stats/tests/test_forest.py b/treeple/stats/tests/test_forest.py
@@ -238,7 +238,7 @@ def test_comight_repeated_feature_sets(seed):
 
 @pytest.mark.parametrize(
     ("use_bottleneck", "use_sparse"),
-    itertools.product([False, True], [False, True]),
+    itertools.product([True, False], [True, False]),
 )
 def test_build_coleman_forest(use_bottleneck: bool, use_sparse: bool):
     """Simple test for building a Coleman forest.

diff --git a/treeple/stats/tests/test_utils.py b/treeple/stats/tests/test_utils.py
@@ -67,3 +67,24 @@ def test_non_nan_samples(use_bottleneck: bool):
     expected = np.array([0, 2])
     actual = utils._non_nan_samples(posterior_array)
     np.testing.assert_array_equal(expected, actual)
+
+
+@pytest.mark.parametrize("use_bottleneck", [True, False])
+def test_nanmean_f(use_bottleneck: bool):
+    if use_bottleneck and utils.DISABLE_BN_ENV_VAR in os.environ:
+        del os.environ[utils.DISABLE_BN_ENV_VAR]
+        importlib.reload(utils)
+    else:
+        os.environ[utils.DISABLE_BN_ENV_VAR] = "1"
+        importlib.reload(utils)
+
+    posterior_array = np.array(
+        [
+            [1, 2, np.nan],
+            [3, 4, np.nan],
+        ]
+    )
+
+    expected = np.array([1.5, 3.5])
+    actual = utils.nanmean_f(posterior_array, axis=1)
+    np.testing.assert_array_equal(expected, actual)