TST Use _sklearn_version rather than version in pickle tests (scikit-learn#27554)

lesteve · web-flow · commit 19932fcdbf79 · 2023-10-09T22:46:49.000+02:00
diff --git a/sklearn/tests/test_base.py b/sklearn/tests/test_base.py
@@ -425,7 +425,7 @@ def test_pickle_version_warning_is_not_raised_with_matching_version():
     iris = datasets.load_iris()
     tree = DecisionTreeClassifier().fit(iris.data, iris.target)
     tree_pickle = pickle.dumps(tree)
-    assert b"version" in tree_pickle
+    assert b"_sklearn_version" in tree_pickle
     tree_restored = assert_no_warnings(pickle.loads, tree_pickle)
 
     # test that we can predict with the restored decision tree classifier
@@ -478,7 +478,7 @@ def test_pickle_version_warning_is_issued_when_no_version_info_in_pickle():
     tree = TreeNoVersion().fit(iris.data, iris.target)
 
     tree_pickle_noversion = pickle.dumps(tree)
-    assert b"version" not in tree_pickle_noversion
+    assert b"_sklearn_version" not in tree_pickle_noversion
     message = pickle_error_message.format(
         estimator="TreeNoVersion",
         old_version="pre-0.18",
diff --git a/sklearn/utils/estimator_checks.py b/sklearn/utils/estimator_checks.py
@@ -2078,7 +2078,7 @@ def check_estimators_pickle(name, estimator_orig, readonly_memmap=False):
         ):
             # strict check for sklearn estimators that are not implemented in test
             # modules.
-            assert b"version" in pickled_estimator
+            assert b"_sklearn_version" in pickled_estimator
         unpickled_estimator = pickle.loads(pickled_estimator)
 
     result = dict()
