1  import pytest
 2  
 3  import mlflow
 4  from mlflow.exceptions import MlflowException
 5  
 6  
 7  class UnpicklableModel(mlflow.pyfunc.PythonModel):
 8      def __init__(self, path):
 9          with open(path, "w+") as f:
10              pass
11  
12          self.not_a_file = f
13  
14  
15  def test_pyfunc_unpicklable_exception(tmp_path):
16      model = UnpicklableModel(tmp_path / "model.pkl")
17  
18      with pytest.raises(
19          MlflowException,
20          match="Please save the model into a python file and use code-based logging method instead",
21      ):
22          mlflow.pyfunc.save_model(python_model=model, path=tmp_path / "model")