/ tests / deployments / test_utils.py
test_utils.py
 1  import pytest
 2  
 3  from mlflow.deployments.utils import (
 4      get_deployments_target,
 5      set_deployments_target,
 6  )
 7  from mlflow.exceptions import MlflowException
 8  
 9  
10  def test_set_deployments_target(monkeypatch):
11      monkeypatch.setattr("mlflow.deployments.utils._deployments_target", None)
12  
13      valid_target = "databricks"
14      set_deployments_target(valid_target)
15      assert get_deployments_target() == valid_target
16  
17      valid_uri = "http://localhost"
18      set_deployments_target(valid_uri)
19      assert get_deployments_target() == valid_uri
20  
21      invalid_uri = "localhost"
22      with pytest.raises(
23          MlflowException, match="The target provided is not a valid uri or 'databricks'"
24      ):
25          set_deployments_target(invalid_uri)
26  
27  
28  def test_get_deployments_target(monkeypatch):
29      monkeypatch.setattr("mlflow.deployments.utils._deployments_target", None)
30      monkeypatch.delenv("MLFLOW_DEPLOYMENTS_TARGET", raising=False)
31  
32      with pytest.raises(MlflowException, match="No deployments target has been set"):
33          get_deployments_target()
34  
35      valid_uri = "http://localhost"
36      monkeypatch.setattr("mlflow.deployments.utils._deployments_target", valid_uri)
37      assert get_deployments_target() == valid_uri
38  
39      monkeypatch.delenv("MLFLOW_DEPLOYMENTS_TARGET", raising=False)
40      set_deployments_target(valid_uri)
41      assert get_deployments_target() == valid_uri