Files
2026-07-13 13:22:34 +08:00

134 lines
4.8 KiB
Python

import pytest
from mlflow.entities.trace_metrics import (
AggregationType,
MetricAggregation,
MetricDataPoint,
MetricViewType,
)
from mlflow.protos import service_pb2 as pb
@pytest.mark.parametrize(
("view_type", "expected_proto"),
zip(MetricViewType, pb.MetricViewType.values(), strict=True),
)
def test_trace_metrics_view_type(view_type: MetricViewType, expected_proto: pb.MetricViewType):
assert view_type.to_proto() == expected_proto
@pytest.mark.parametrize(
("aggregation_type", "expected_proto"),
zip(AggregationType, pb.AggregationType.values(), strict=True),
)
def test_trace_metrics_aggregation_type_to_proto(
aggregation_type: AggregationType, expected_proto: pb.AggregationType
):
assert aggregation_type.to_proto() == expected_proto
def test_metrics_aggregation_to_proto_without_percentile():
aggregation = MetricAggregation(aggregation_type=AggregationType.AVG)
proto = aggregation.to_proto()
assert proto.aggregation_type == pb.AggregationType.AVG
assert not proto.HasField("percentile_value")
@pytest.mark.parametrize(
("percentile_value"),
[50.0, 75.0, 90.0, 95.0, 99.0, 99.9],
)
def test_metrics_aggregation_percentile_values(percentile_value: float):
aggregation = MetricAggregation(
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
)
proto = aggregation.to_proto()
assert proto.percentile_value == percentile_value
def test_metrics_aggregation_percentile_requires_value():
with pytest.raises(ValueError, match="Percentile value is required for PERCENTILE aggregation"):
MetricAggregation(aggregation_type=AggregationType.PERCENTILE)
@pytest.mark.parametrize("percentile_value", [-1.0, -0.1, 100.1, 101.0, 1000.0])
def test_metrics_aggregation_percentile_value_out_of_range(percentile_value: float):
with pytest.raises(ValueError, match="Percentile value must be between 0 and 100"):
MetricAggregation(
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
)
@pytest.mark.parametrize("percentile_value", [0.0, 0.1, 50.0, 99.9, 100.0])
def test_metrics_aggregation_percentile_value_valid_range(percentile_value: float):
aggregation = MetricAggregation(
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
)
assert aggregation.percentile_value == percentile_value
@pytest.mark.parametrize(
"agg_type",
[t for t in AggregationType if t is not AggregationType.PERCENTILE],
)
def test_metrics_aggregation_non_percentile_with_value_raises(agg_type: AggregationType):
with pytest.raises(
ValueError, match="Percentile value is only allowed for PERCENTILE aggregation"
):
MetricAggregation(aggregation_type=agg_type, percentile_value=50.0)
def test_trace_metrics_metric_data_point_from_proto():
metric_data_point_proto = pb.MetricDataPoint(
metric_name="latency",
dimensions={"status": "OK"},
values={"avg": 150.5, "p99": 200},
)
assert MetricDataPoint.from_proto(metric_data_point_proto) == MetricDataPoint(
metric_name="latency",
dimensions={"status": "OK"},
values={"avg": 150.5, "p99": 200},
)
def test_trace_metrics_metric_data_point_to_proto():
metric_data_point = MetricDataPoint(
metric_name="latency",
dimensions={"status": "OK", "model": "gpt-4"},
values={"avg": 150.5, "p99": 200.0},
)
proto = metric_data_point.to_proto()
assert proto.metric_name == "latency"
assert dict(proto.dimensions) == {"status": "OK", "model": "gpt-4"}
assert dict(proto.values) == {"avg": 150.5, "p99": 200.0}
@pytest.mark.parametrize(
("view_type", "expected_proto"),
zip(MetricViewType, pb.MetricViewType.values(), strict=True),
)
def test_trace_metrics_view_type_from_proto(view_type: MetricViewType, expected_proto: int):
assert MetricViewType.from_proto(expected_proto) == view_type
@pytest.mark.parametrize(
"agg_type",
[t for t in AggregationType if t is not AggregationType.PERCENTILE],
)
def test_metrics_aggregation_from_proto_without_percentile(agg_type: AggregationType):
proto = pb.MetricAggregation(aggregation_type=agg_type.to_proto())
aggregation = MetricAggregation.from_proto(proto)
assert aggregation.aggregation_type == agg_type
assert aggregation.percentile_value is None
@pytest.mark.parametrize("percentile_value", [50.0, 75.0, 90.0, 95.0, 99.0, 99.9])
def test_metrics_aggregation_from_proto_with_percentile(percentile_value: float):
proto = pb.MetricAggregation(
aggregation_type=pb.AggregationType.PERCENTILE,
percentile_value=percentile_value,
)
aggregation = MetricAggregation.from_proto(proto)
assert aggregation.aggregation_type == AggregationType.PERCENTILE
assert aggregation.percentile_value == percentile_value