fix tests

arogozhnikov · Mar 5, 2020 · 39dd942 · 39dd942
1 parent 718c7d2
commit 39dd942
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 1 deletion.
diff --git a/hep_ml/gradientboosting.py b/hep_ml/gradientboosting.py
@@ -149,6 +149,7 @@ def _get_train_features(self, X):
         if self.train_features is None:
             return X
         else:
+            assert all(f in X.columns for f in self.train_features), [X.columns, self.train_features]
             return X.loc[:, self.train_features]
 
     def staged_decision_function(self, X):

diff --git a/tests/test_gradientboosting.py b/tests/test_gradientboosting.py
@@ -61,6 +61,7 @@ def test_gradient_boosting(n_samples=1000, distance = 0.6):
         assert result >= 0.7, "The quality is too poor: {} with loss: {}".format(result, loss)
 
     trainX['fake_request'] = numpy.random.randint(0, 4, size=len(trainX))
+    testX['fake_request'] = numpy.random.randint(0, 4, size=len(testX))
     for loss in [losses.MSELossFunction(),
                  losses.MAELossFunction(),
                  losses.RankBoostLossFunction(request_column='fake_request')]:

diff --git a/tests/test_metrics.py b/tests/test_metrics.py
@@ -244,7 +244,7 @@ def test_metrics_clear(n_samples=2000, knn=50, uniform_class=0):
     features = X.columns[:1]
 
     mask = (y == uniform_class)
-    X_clear = X.ix[mask, :]
+    X_clear = X[mask]
     y_clear = y[mask]
     sample_weight_clear = sample_weight[mask]
     predictions_clear = predictions[mask]