1  import json
  2  import os
  3  from unittest import mock
  4  
  5  import pytest
  6  from click.testing import CliRunner
  7  
  8  from mlflow.deployments import cli
  9  from mlflow.exceptions import MlflowException
 10  
 11  f_model_uri = "fake_model_uri"
 12  f_name = "fake_deployment_name"
 13  f_flavor = "fake_flavor"
 14  f_target = "faketarget"
 15  runner = CliRunner()
 16  
 17  
 18  def test_create():
 19      res = runner.invoke(
 20          cli.create_deployment,
 21          ["--flavor", f_flavor, "--model-uri", f_model_uri, "--target", f_target, "--name", f_name],
 22      )
 23      assert f"{f_flavor} deployment {f_name} is created" in res.stdout
 24      res = runner.invoke(
 25          cli.create_deployment, ["-f", f_flavor, "-m", f_model_uri, "-t", f_target, "--name", f_name]
 26      )
 27      assert f"{f_flavor} deployment {f_name} is created" in res.stdout
 28  
 29  
 30  def test_update():
 31      res = runner.invoke(
 32          cli.update_deployment,
 33          ["--flavor", f_flavor, "--model-uri", f_model_uri, "--target", f_target, "--name", f_name],
 34      )
 35      assert f"Deployment {f_name} is updated (with flavor {f_flavor})" in res.stdout
 36  
 37  
 38  def test_delete():
 39      res = runner.invoke(cli.delete_deployment, ["--name", f_name, "--target", f_target])
 40      assert f"Deployment {f_name} is deleted" in res.stdout
 41  
 42  
 43  def test_update_no_flavor():
 44      res = runner.invoke(
 45          cli.update_deployment, ["--name", f_name, "--target", f_target, "-m", f_model_uri]
 46      )
 47      assert f"Deployment {f_name} is updated (with flavor None)" in res.stdout
 48  
 49  
 50  def test_list():
 51      res = runner.invoke(cli.list_deployment, ["--target", f_target])
 52      assert f"[{{'name': '{f_name}'}}]" in res.stdout
 53  
 54  
 55  def test_create_deployment_with_custom_args():
 56      res = runner.invoke(
 57          cli.create_deployment,
 58          [
 59              "--model-uri",
 60              f_model_uri,
 61              "--target",
 62              f_target,
 63              "--name",
 64              f_name,
 65              "-C",
 66              "raiseError=True",
 67          ],
 68      )
 69      assert isinstance(res.exception, RuntimeError)
 70  
 71  
 72  def test_delete_deployment_with_custom_args():
 73      res = runner.invoke(
 74          cli.delete_deployment,
 75          ["--target", f_target, "--name", f_name, "-C", "raiseError=True"],
 76      )
 77      assert isinstance(res.exception, RuntimeError)
 78  
 79  
 80  def test_get():
 81      res = runner.invoke(cli.get_deployment, ["--name", f_name, "--target", f_target])
 82      assert "key1: val1" in res.stdout
 83      assert "key2: val2" in res.stdout
 84  
 85  
 86  @pytest.mark.skipif(
 87      "MLFLOW_SKINNY" in os.environ,
 88      reason="Skinny Client does not support predict due to the pandas dependency",
 89  )
 90  def test_predict(tmp_path):
 91      temp_input_file_path = tmp_path.joinpath("input.json")
 92      temp_input_file_path.write_text('{"data": [5000]}')
 93  
 94      temp_output_file_path = tmp_path.joinpath("output.json")
 95      res = runner.invoke(
 96          cli.predict, ["--target", f_target, "--name", f_name, "--input-path", temp_input_file_path]
 97      )
 98      assert '{"predictions": [1, 2, 3]}' in res.stdout
 99  
100      res = runner.invoke(
101          cli.predict,
102          [
103              "--target",
104              f_target,
105              "--name",
106              f_name,
107              "--input-path",
108              temp_input_file_path,
109              "--output-path",
110              temp_output_file_path,
111          ],
112      )
113      with open(temp_output_file_path) as f:
114          assert json.load(f) == {"predictions": [1, 2, 3]}
115  
116  
117  def test_target_help():
118      res = runner.invoke(cli.target_help, ["--target", f_target])
119      assert "Target help is called" in res.stdout
120  
121  
122  def test_run_local():
123      res = runner.invoke(
124          cli.run_local, ["-f", f_flavor, "-m", f_model_uri, "-t", f_target, "--name", f_name]
125      )
126      assert f"Deployed locally at the key {f_name}" in res.stdout
127      assert f"using the model from {f_model_uri}." in res.stdout
128      assert f"It's flavor is {f_flavor} and config is {{}}" in res.stdout
129  
130  
131  @pytest.mark.skipif(
132      "MLFLOW_SKINNY" in os.environ,
133      reason="Skinny Client does not support explain due to the pandas dependency",
134  )
135  def test_explain(tmp_path):
136      temp_input_file_path = tmp_path.joinpath("input.json")
137      temp_input_file_path.write_text('{"data": [5000]}')
138      res = runner.invoke(
139          cli.explain, ["--target", f_target, "--name", f_name, "--input-path", temp_input_file_path]
140      )
141      assert "1" in res.stdout
142  
143  
144  def test_explain_with_no_target_implementation(tmp_path):
145      file_path = tmp_path.joinpath("input.json")
146      file_path.write_text('{"data": [5000]}')
147      mock_error = MlflowException("MOCK ERROR")
148      with mock.patch.object(CliRunner, "invoke", return_value=mock_error) as mock_explain:
149          res = runner.invoke(
150              cli.explain, ["--target", f_target, "--name", f_name, "--input-path", file_path]
151          )
152          assert type(res) == MlflowException
153          mock_explain.assert_called_once()