1  from pathlib import Path
 2  
 3  from clint.config import Config
 4  from clint.index import SymbolIndex
 5  from clint.linter import Position, Range, lint_file
 6  from clint.rules.forbidden_set_active_model_usage import ForbiddenSetActiveModelUsage
 7  
 8  
 9  def test_forbidden_set_active_model_usage(index: SymbolIndex) -> None:
10      code = """
11  import mlflow
12  
13  # Bad
14  mlflow.set_active_model("model_name")
15  
16  # Good
17  mlflow._set_active_model("model_name")
18  
19  # Bad - with aliasing
20  from mlflow import set_active_model
21  set_active_model("model_name")
22  
23  # Good - with aliasing
24  from mlflow import _set_active_model
25  _set_active_model("model_name")
26  """
27      config = Config(select={ForbiddenSetActiveModelUsage.name})
28      violations = lint_file(Path("test.py"), code, config, index)
29      assert len(violations) == 3
30      assert all(isinstance(v.rule, ForbiddenSetActiveModelUsage) for v in violations)
31      assert violations[0].range == Range(Position(4, 0))  # mlflow.set_active_model call
32      assert violations[1].range == Range(Position(10, 0))  # from mlflow import set_active_model
33      assert violations[2].range == Range(Position(11, 0))  # direct set_active_model call