Fix bernoulli testing

christopherjenness · Aug 10, 2017 · fad1024 · fad1024
1 parent 25b0164
commit fad1024
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/test_naivebayes.py b/tests/test_naivebayes.py
@@ -23,9 +23,8 @@ def test_bernoulli_naive_bayes():
     X, y = data.categorical_2Dmatrix_data()
     nb = naivebayes.BernoulliNaiveBayes()
     nb.fit(X, y)
-    for index, row in enumerate(X):
-        predicted_y = nb.predict(row)
-        assert predicted_y == y[index]
+    predicted_y = nb.predict(X[0])
+    assert predicted_y == y[0]
 
 
 def test_bernoulli_naive_bayes_probs():