wehub-resource-sync
87 lines
2.4 KiB
Python
87 lines
2.4 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from clint.config import Config
|
|
from clint.index import SymbolIndex
|
|
from clint.linter import Position, Range, lint_file
|
|
from clint.rules import GetArtifactUri
|
|
|
|
|
|
def test_get_artifact_uri_in_rst_example(index: SymbolIndex) -> None:
|
|
code = """
|
|
Documentation
|
|
=============
|
|
|
|
Here's an example:
|
|
|
|
.. code-block:: python
|
|
|
|
import mlflow
|
|
|
|
with mlflow.start_run():
|
|
mlflow.sklearn.log_model(model, "model")
|
|
model_uri = mlflow.get_artifact_uri("model")
|
|
print(model_uri)
|
|
"""
|
|
config = Config(select={GetArtifactUri.name}, example_rules=[GetArtifactUri.name])
|
|
violations = lint_file(Path("test.rst"), code, config, index)
|
|
assert len(violations) == 1
|
|
assert violations[0].rule.name == GetArtifactUri.name
|
|
assert violations[0].range == Range(Position(12, 20))
|
|
|
|
|
|
@pytest.mark.parametrize("suffix", [".md", ".mdx"])
|
|
def test_get_artifact_uri_in_markdown_example(index: SymbolIndex, suffix: str) -> None:
|
|
code = """
|
|
# Documentation
|
|
|
|
Here's an example:
|
|
|
|
```python
|
|
import mlflow
|
|
|
|
with mlflow.start_run():
|
|
mlflow.sklearn.log_model(model, "model")
|
|
model_uri = mlflow.get_artifact_uri("model")
|
|
print(model_uri)
|
|
```
|
|
"""
|
|
config = Config(select={GetArtifactUri.name}, example_rules=[GetArtifactUri.name])
|
|
violations = lint_file(Path("test").with_suffix(suffix), code, config, index)
|
|
assert len(violations) == 1
|
|
assert violations[0].rule.name == GetArtifactUri.name
|
|
assert violations[0].range == Range(Position(10, 16))
|
|
|
|
|
|
def test_get_artifact_uri_not_in_regular_python_files(index: SymbolIndex) -> None:
|
|
code = """
|
|
import mlflow
|
|
|
|
with mlflow.start_run():
|
|
model_uri = mlflow.get_artifact_uri("model")
|
|
print(model_uri)
|
|
"""
|
|
config = Config(select={GetArtifactUri.name}, example_rules=[GetArtifactUri.name])
|
|
violations = lint_file(Path("test.py"), code, config, index)
|
|
assert len(violations) == 0
|
|
|
|
|
|
def test_get_artifact_uri_without_log_model_allowed(index: SymbolIndex) -> None:
|
|
code = """
|
|
Documentation
|
|
=============
|
|
|
|
Here's an example:
|
|
|
|
.. code-block:: python
|
|
|
|
import mlflow
|
|
|
|
# This should be allowed - no log_model in the example
|
|
model_uri = mlflow.get_artifact_uri("some_model")
|
|
loaded_model = mlflow.sklearn.load_model(model_uri)
|
|
"""
|
|
config = Config(select={GetArtifactUri.name}, example_rules=[GetArtifactUri.name])
|
|
violations = lint_file(Path("test.rst"), code, config, index)
|
|
assert len(violations) == 0
|