Add unit test for rank checks

zillow · Sep 4, 2024 · e6d7d63 · e6d7d63
1 parent 7081e51
commit e6d7d63
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/quantile_forest/tests/test_quantile_forest.py b/quantile_forest/tests/test_quantile_forest.py
@@ -682,6 +682,10 @@ def check_quantile_ranks(name):
     assert np.all(y_ranks <= 1)
     assert np.all(y_ranks[:, 0] < y_ranks[:, 1])
 
+    # Check error if training and test number of targets are not equal.
+    est.fit(X_train, y_train[:, 0])  # training target size = 1
+    assert_raises(ValueError, est.quantile_ranks, X_test, y_test[:, :2])  # test target size = 2
+
 
 @pytest.mark.parametrize("name", FOREST_REGRESSORS)
 def test_quantile_ranks(name):