Write new tests for model summaries

AdamPSU · Sep 1, 2024 · 651e649 · 651e649
1 parent 822f2db
commit 651e649
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/kavian/tables/linear_model.py b/kavian/tables/linear_model.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from rich.panel import Panel
-from kavian.tables.base import BaseRegressorSummary, BaseClassifierSummary
+from kavian.tables.base import BaseRegressorSummary, BaseBinaryClassifierSummary
 from kavian.tables.utils import format_stat
 
 class RegularizedRegressorSummary(BaseRegressorSummary):
@@ -46,7 +46,7 @@ def sparse_coefficients(self):
         return zeros
 
 
-class BinaryClassifierSummary(BaseClassifierSummary):
+class BinaryClassifierSummary(BaseBinaryClassifierSummary):
     def __init__(self, estimator, X, y):
         super().__init__(estimator, X, y)
 

diff --git a/kavian/tests/test_summaries/test_base.py b/kavian/tests/test_summaries/test_base.py
@@ -7,7 +7,6 @@
 
 from kavian.tables.base import (
     BaseRegressorSummary,
-    BaseClassifierSummary,
     SimpleRegressorSummary,
     SimpleClassifierSummary
 )
@@ -53,10 +52,11 @@ def test_compatibility(get_diabetes, get_breast_cancer):
 
 def test_too_many_entries():
     with (pytest.raises(ValueError)):
-        test = TooManyEntriesSummary(LinearRegression(), X_RANDOM, Y_RANDOM)
-        test.summary()
-
+        model = LinearRegression()
+        model.fit(X_RANDOM, Y_RANDOM)
 
+        test = TooManyEntriesSummary(model, X_RANDOM, Y_RANDOM)
+        test.summary()