diff --git a/src/documents/tests/data/v1.17.4.model.pickle b/src/documents/tests/data/v1.17.4.model.pickle deleted file mode 100644 index 4b2734607..000000000 Binary files a/src/documents/tests/data/v1.17.4.model.pickle and /dev/null differ diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py index a1bb9baa5..5e8195166 100644 --- a/src/documents/tests/test_classifier.py +++ b/src/documents/tests/test_classifier.py @@ -1,4 +1,3 @@ -import pickle import re import warnings from pathlib import Path @@ -447,18 +446,28 @@ class TestClassifier(DirectoriesMixin, TestCase): self.generate_train_and_save() - real_loads = pickle.loads + fake_warning = warnings.WarningMessage( + message=InconsistentVersionWarning( + estimator_name="MLPClassifier", + current_sklearn_version="1.0", + original_sklearn_version="0.9", + ), + category=InconsistentVersionWarning, + filename="", + lineno=0, + ) - def loads_with_warning(data): - warnings.warn( - "Trying to unpickle estimator from version 0.0 when using version 1.0.", - InconsistentVersionWarning, - ) - return real_loads(data) + real_catch_warnings = warnings.catch_warnings + + class PatchedCatchWarnings(real_catch_warnings): + def __enter__(self): + w = super().__enter__() + w.append(fake_warning) + return w with mock.patch( - "documents.classifier.pickle.loads", - side_effect=loads_with_warning, + "documents.classifier.warnings.catch_warnings", + PatchedCatchWarnings, ): with self.assertRaises(IncompatibleClassifierVersionError): self.classifier.load()