wehub-resource-sync
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
from pathlib import Path
|
|
|
|
from clint.config import Config
|
|
from clint.index import SymbolIndex
|
|
from clint.linter import Position, Range, lint_file
|
|
from clint.rules.log_model_artifact_path import LogModelArtifactPath
|
|
|
|
|
|
def test_log_model_artifact_path(index: SymbolIndex) -> None:
|
|
code = """
|
|
import mlflow
|
|
|
|
# Bad - using deprecated artifact_path positionally
|
|
mlflow.sklearn.log_model(model, "model")
|
|
|
|
# Bad - using deprecated artifact_path as keyword
|
|
mlflow.tensorflow.log_model(model, artifact_path="tf_model")
|
|
|
|
# Good - using the new 'name' parameter
|
|
mlflow.sklearn.log_model(model, name="my_model")
|
|
|
|
# Good - spark flavor is exempted from this rule
|
|
mlflow.spark.log_model(spark_model, "spark_model")
|
|
|
|
# Bad - another flavor with artifact_path
|
|
mlflow.pytorch.log_model(model, artifact_path="pytorch_model")
|
|
"""
|
|
config = Config(select={LogModelArtifactPath.name})
|
|
violations = lint_file(Path("test.py"), code, config, index)
|
|
assert len(violations) == 3
|
|
assert all(isinstance(v.rule, LogModelArtifactPath) for v in violations)
|
|
assert violations[0].range == Range(Position(4, 0))
|
|
assert violations[1].range == Range(Position(7, 0))
|
|
assert violations[2].range == Range(Position(16, 0))
|