1  from pathlib import Path
 2  
 3  import pytest
 4  from clint.config import Config
 5  from clint.index import SymbolIndex
 6  from clint.linter import Position, Range, lint_file
 7  from clint.rules.unknown_mlflow_function import UnknownMlflowFunction
 8  
 9  
10  def test_unknown_mlflow_function(index: SymbolIndex) -> None:
11      code = '''
12  def bad():
13      """
14      .. code-block:: python
15  
16          import mlflow
17  
18          mlflow.foo()
19  
20      """
21  
22  
23  def good():
24  
25      """
26      .. code-block:: python
27  
28          import mlflow
29  
30          mlflow.log_param("k", "v")
31  
32      """
33  '''
34      config = Config(select={UnknownMlflowFunction.name}, example_rules=[UnknownMlflowFunction.name])
35      violations = lint_file(Path("test.py"), code, config, index)
36      assert len(violations) == 1
37      assert all(isinstance(v.rule, UnknownMlflowFunction) for v in violations)
38      assert violations[0].range == Range(Position(7, 8))
39  
40  
41  @pytest.mark.parametrize("suffix", [".md", ".mdx"])
42  def test_unknown_mlflow_function_markdown(index: SymbolIndex, suffix: str) -> None:
43      code = """
44  # Bad
45  
46  ```python
47  import mlflow
48  
49  mlflow.foo()
50  ```
51  
52  # Good
53  
54  ```python
55  import mlflow
56  
57  mlflow.log_param("k", "v")
58  ```
59  
60  """
61      config = Config(
62          select={UnknownMlflowFunction.name},
63          example_rules=[UnknownMlflowFunction.name],
64      )
65      violations = lint_file(Path("test").with_suffix(suffix), code, config, index)
66      assert len(violations) == 1
67      assert all(isinstance(v.rule, UnknownMlflowFunction) for v in violations)
68      assert violations[0].range == Range(Position(6, 0))