From fcf11c2c07476851379580ab50b6cac27fc79a0f Mon Sep 17 00:00:00 2001
From: Keith Battocchi <kebatt@microsoft.com>
Date: Fri, 15 May 2020 11:05:20 -0400
Subject: [PATCH] Fix tests

---
 econml/tests/test_dml.py       |  2 +-
 econml/tests/test_drlearner.py |  2 +-
 econml/tests/test_orf.py       | 20 +++++++++++++++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/econml/tests/test_dml.py b/econml/tests/test_dml.py
index 5fcadbe68..5da595070 100644
--- a/econml/tests/test_dml.py
+++ b/econml/tests/test_dml.py
@@ -699,7 +699,7 @@ def test_can_custom_splitter(self):
     def test_can_use_featurizer(self):
         "Test that we can use a featurizer, and that fit is only called during training"
         dml = LinearDMLCateEstimator(LinearRegression(), LinearRegression(),
-                                     fit_cate_intercept=False, featurizer=OneHotEncoder(n_values='auto', sparse=False))
+                                     fit_cate_intercept=False, featurizer=OneHotEncoder(sparse=False))
 
         T = np.tile([1, 2, 3], 6)
         Y = np.array([1, 2, 3, 1, 2, 3])
diff --git a/econml/tests/test_drlearner.py b/econml/tests/test_drlearner.py
index 89f6f2517..baca21165 100644
--- a/econml/tests/test_drlearner.py
+++ b/econml/tests/test_drlearner.py
@@ -713,7 +713,7 @@ def test_sparse(self):
         y_lower, y_upper = sparse_dml.effect_interval(x_test, T0=0, T1=1)
         in_CI = ((y_lower < true_eff) & (true_eff < y_upper))
         # Check that a majority of true effects lie in the 5-95% CI
-        self.assertTrue(in_CI.mean() > 0.8)
+        self.assertGreater(in_CI.mean(), 0.8)
 
     def _test_te(self, learner_instance, tol, te_type="const"):
         if te_type not in ["const", "heterogeneous"]:
diff --git a/econml/tests/test_orf.py b/econml/tests/test_orf.py
index 114062a2b..ada59691c 100644
--- a/econml/tests/test_orf.py
+++ b/econml/tests/test_orf.py
@@ -184,6 +184,20 @@ def test_effect_shape(self):
 
     def test_nuisance_model_has_weights(self):
         """Test whether the correct exception is being raised if model_final doesn't have weights."""
+
+        # Create a wrapper around Lasso that doesn't support weights
+        # since Lasso does natively support them starting in sklearn 0.23
+        class NoWeightModel:
+            def __init__(self):
+                self.model = Lasso()
+
+            def fit(self, X, y):
+                self.model.fit(X, y)
+                return self
+
+            def predict(self, X):
+                return self.model.predict(X)
+
         # Generate data with continuous treatments
         T = np.dot(TestOrthoForest.W[:, TestOrthoForest.support], TestOrthoForest.coefs_T) + \
             TestOrthoForest.eta_sample(TestOrthoForest.n)
@@ -192,14 +206,14 @@ def test_nuisance_model_has_weights(self):
             T * TE + TestOrthoForest.epsilon_sample(TestOrthoForest.n)
         # Instantiate model with most of the default parameters
         est = ContinuousTreatmentOrthoForest(n_jobs=4, n_trees=10,
-                                             model_T=Lasso(),
-                                             model_Y=Lasso())
+                                             model_T=NoWeightModel(),
+                                             model_Y=NoWeightModel())
         est.fit(Y=Y, T=T, X=TestOrthoForest.X, W=TestOrthoForest.W)
         weights_error_msg = (
             "Estimators of type {} do not accept weights. "
             "Consider using the class WeightedModelWrapper from econml.utilities to build a weighted model."
         )
-        self.assertRaisesRegexp(TypeError, weights_error_msg.format("Lasso"),
+        self.assertRaisesRegexp(TypeError, weights_error_msg.format("NoWeightModel"),
                                 est.effect, X=TestOrthoForest.X)
 
     def _test_te(self, learner_instance, expected_te, tol, treatment_type='continuous'):