diff --git a/UnleashClient/impact_metrics/__init__.py b/UnleashClient/impact_metrics/__init__.py new file mode 100644 index 00000000..e11fef60 --- /dev/null +++ b/UnleashClient/impact_metrics/__init__.py @@ -0,0 +1,19 @@ +from .metric_types import ( + CollectedMetric, + Counter, + CounterImpl, + MetricLabels, + MetricOptions, + MetricType, + NumericMetricSample, +) + +__all__ = [ + "CollectedMetric", + "Counter", + "CounterImpl", + "MetricLabels", + "MetricOptions", + "MetricType", + "NumericMetricSample", +] diff --git a/UnleashClient/impact_metrics/metric_types.py b/UnleashClient/impact_metrics/metric_types.py new file mode 100644 index 00000000..6024cc68 --- /dev/null +++ b/UnleashClient/impact_metrics/metric_types.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from threading import RLock +from typing import Any, Dict, List, Optional, Protocol + + +class MetricType(str, Enum): + COUNTER = "counter" + GAUGE = "gauge" + HISTOGRAM = "histogram" + + +MetricLabels = Dict[str, str] + + +@dataclass +class MetricOptions: + name: str + help: str + label_names: List[str] = field(default_factory=list) + + +@dataclass +class NumericMetricSample: + labels: MetricLabels + value: int + + def to_dict(self) -> Dict[str, Any]: + return {"labels": self.labels, "value": self.value} + + +@dataclass +class CollectedMetric: + name: str + help: str + type: MetricType + samples: List[NumericMetricSample] + + def to_dict(self) -> Dict[str, Any]: + return { + "name": self.name, + "help": self.help, + "type": self.type.value, + "samples": [s.to_dict() for s in self.samples], + } + + +def _get_label_key(labels: Optional[MetricLabels]) -> str: + if not labels: + return "" + return ",".join(f"{k}={v}" for k, v in sorted(labels.items())) + + +def _parse_label_key(key: str) -> MetricLabels: + if not key: + return {} + labels: MetricLabels = {} + for pair in key.split(","): + if "=" in pair: + k, v = pair.split("=", 1) + labels[k] = v + return labels + + +class Counter(Protocol): + def inc(self, value: int = 1, labels: Optional[MetricLabels] = None) -> None: ... + + +class CounterImpl: + def __init__(self, opts: MetricOptions) -> None: + self._opts = opts + self._values: Dict[str, int] = {} + self._lock = RLock() + + def inc(self, value: int = 1, labels: Optional[MetricLabels] = None) -> None: + key = _get_label_key(labels) + with self._lock: + current = self._values.get(key, 0) + self._values[key] = current + value + + def collect(self) -> CollectedMetric: + samples: List[NumericMetricSample] = [] + + with self._lock: + for key in list(self._values.keys()): + value = self._values.pop(key) + samples.append( + NumericMetricSample(labels=_parse_label_key(key), value=value) + ) + + if not samples: + samples.append(NumericMetricSample(labels={}, value=0)) + + return CollectedMetric( + name=self._opts.name, + help=self._opts.help, + type=MetricType.COUNTER, + samples=samples, + ) diff --git a/tests/unit_tests/impact_metrics/__init__.py b/tests/unit_tests/impact_metrics/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/impact_metrics/test_metric_types.py b/tests/unit_tests/impact_metrics/test_metric_types.py new file mode 100644 index 00000000..1ab6444d --- /dev/null +++ b/tests/unit_tests/impact_metrics/test_metric_types.py @@ -0,0 +1,76 @@ +from UnleashClient.impact_metrics.metric_types import ( + CounterImpl, + MetricOptions, +) + + +def test_counter_increments_by_default_value(): + counter = CounterImpl(MetricOptions(name="test_counter", help="testing")) + + counter.inc() + + result = counter.collect() + + assert result.to_dict() == { + "name": "test_counter", + "help": "testing", + "type": "counter", + "samples": [{"labels": {}, "value": 1}], + } + + +def test_counter_increments_with_custom_value_and_labels(): + counter = CounterImpl(MetricOptions(name="labeled_counter", help="with labels")) + + counter.inc(3, {"foo": "bar"}) + counter.inc(2, {"foo": "bar"}) + + result = counter.collect() + + assert result.to_dict() == { + "name": "labeled_counter", + "help": "with labels", + "type": "counter", + "samples": [{"labels": {"foo": "bar"}, "value": 5}], + } + + +def test_different_label_combinations_are_stored_separately(): + counter = CounterImpl(MetricOptions(name="multi_label", help="label test")) + + counter.inc(1, {"a": "x"}) + counter.inc(2, {"b": "y"}) + counter.inc(3) + + result = counter.collect() + + result.samples.sort(key=lambda s: s.value) + + assert result.to_dict() == { + "name": "multi_label", + "help": "label test", + "type": "counter", + "samples": [ + {"labels": {"a": "x"}, "value": 1}, + {"labels": {"b": "y"}, "value": 2}, + {"labels": {}, "value": 3}, + ], + } + + +def test_collect_returns_counter_with_zero_value_after_flushing_previous_values(): + counter = CounterImpl(MetricOptions(name="flush_test", help="flush")) + + counter.inc(1) + first = counter.collect() + assert first is not None + assert len(first.samples) == 1 + + second = counter.collect() + + assert second.to_dict() == { + "name": "flush_test", + "help": "flush", + "type": "counter", + "samples": [{"labels": {}, "value": 0}], + }