refactor metrics to remove trackers dict

logprep/metrics/metrics.py

@@ -1,5 +1,6 @@
 """This module tracks, calculates, exposes and resets logprep metrics"""
 from abc import ABC, abstractmethod
+from typing import Union
 
 from attr import define, field, validators
 from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram
@@ -21,47 +22,50 @@ class Metric(ABC):
         ],
         default={},
     )
-    trackers: dict = None
     _registry: CollectorRegistry = field(default=None)
     _prefix: str = field(default="logprep_")
+    tracker: Union[Counter, Histogram, Gauge] = field(init=False, default=None)
 
     @property
     def fullname(self):
         """returns the fullname"""
         return f"{self._prefix}{self.name}"
 
-    def init_tracker(self):
+    def init_tracker(self) -> None:
         """initializes the tracker and adds it to the trackers dict"""
-        tracker = None
         try:
             if isinstance(self, CounterMetric):
-                tracker = Counter(
+                self.tracker = Counter(
                     name=self.fullname,
                     documentation=self.description,
                     labelnames=self.labels.keys(),
                     registry=self._registry,
                 )
             if isinstance(self, HistogramMetric):
-                tracker = Histogram(
+                self.tracker = Histogram(
                     name=self.fullname,
                     documentation=self.description,
                     labelnames=self.labels.keys(),
                     buckets=(0.000001, 0.00001, 0.0001, 0.001, 0.01, 0.1, 1),
                     registry=self._registry,
                 )
             if isinstance(self, GaugeMetric):
-                tracker = Gauge(
+                self.tracker = Gauge(
                     name=self.fullname,
                     documentation=self.description,
                     labelnames=self.labels.keys(),
                     registry=self._registry,
                     multiprocess_mode="liveall",
                 )
-            tracker.labels(**self.labels)
-
-            self.trackers.update({self.fullname: tracker})
-        except ValueError:
-            self.trackers.get(self.fullname).labels(**self.labels)
+        except ValueError as error:
+            # pylint: disable=protected-access
+            self.tracker = self._registry._names_to_collectors.get(self.fullname)
+            # pylint: enable=protected-access
+            if not isinstance(self.tracker, METRIC_TO_COLLECTOR_TYPE[type(self)]):
+                raise ValueError(
+                    f"Metric {self.fullname} already exists with different type"
+                ) from error
+        self.tracker.labels(**self.labels)
 
     @abstractmethod
     def __add__(self, other):
@@ -72,40 +76,41 @@ def __add__(self, other):
 class CounterMetric(Metric):
     """Wrapper for prometheus Counter metric"""
 
-    trackers: dict = {}
-
     def __add__(self, other):
         return self.add_with_labels(other, self.labels)
 
     def add_with_labels(self, other, labels):
         """Add with labels"""
         labels = self.labels | labels
-        self.trackers.get(self.fullname).labels(**labels).inc(other)
+        self.tracker.labels(**labels).inc(other)
         return self
 
 
 @define(kw_only=True)
 class HistogramMetric(Metric):
     """Wrapper for prometheus Histogram metric"""
 
-    trackers: dict = {}
-
     def __add__(self, other):
-        self.trackers.get(self.fullname).labels(**self.labels).observe(other)
+        self.tracker.labels(**self.labels).observe(other)
         return self
 
 
 @define(kw_only=True)
 class GaugeMetric(Metric):
     """Wrapper for prometheus Gauge metric""" ""
 
-    trackers: dict = {}
-
     def __add__(self, other):
         return self.add_with_labels(other, self.labels)
 
     def add_with_labels(self, other, labels):
         """Add with labels"""
         labels = self.labels | labels
-        self.trackers.get(self.fullname).labels(**labels).set(other)
+        self.tracker.labels(**labels).set(other)
         return self
+
+
+METRIC_TO_COLLECTOR_TYPE = {
+    CounterMetric: Counter,
+    HistogramMetric: Histogram,
+    GaugeMetric: Gauge,
+}

tests/unit/metrics/test_metrics.py

@@ -4,29 +4,35 @@
 
 import re
 
+import pytest
 from prometheus_client import CollectorRegistry, Counter, Histogram, generate_latest
 
-from logprep.metrics.metrics import CounterMetric, HistogramMetric, GaugeMetric
+from logprep.metrics.metrics import CounterMetric, GaugeMetric, HistogramMetric
 
 
 class TestsMetrics:
     def setup_method(self):
         self.custom_registry = CollectorRegistry()
 
-    def teardown_method(self):
-        CounterMetric(name="", description="").trackers.clear()
-        HistogramMetric(name="", description="").trackers.clear()
-        GaugeMetric(name="", description="").trackers.clear()
-
-    def test_init_tracker_creates_metric(self):
+    def test_init_tracker_returns_collector(self):
         metric = CounterMetric(
             name="testmetric",
             description="empty description",
             labels={"A": "a"},
             registry=self.custom_registry,
         )
         metric.init_tracker()
-        assert isinstance(metric.trackers.get(metric.fullname), Counter)
+        assert isinstance(metric.tracker, Counter)
+
+    def test_init_tracker_does_not_raise_if_initialized_twice(self):
+        metric = CounterMetric(
+            name="testmetric",
+            description="empty description",
+            labels={"A": "a"},
+            registry=self.custom_registry,
+        )
+        assert metric.init_tracker() == metric.init_tracker()
+        assert isinstance(metric.tracker, Counter)
 
     def test_counter_metric_sets_labels(self):
         metric = CounterMetric(
@@ -36,8 +42,7 @@ def test_counter_metric_sets_labels(self):
             registry=self.custom_registry,
         )
         metric.init_tracker()
-        assert metric.trackers.get(metric.fullname)._labelnames == ("pipeline",)
-        assert ("pipeline-1",) in metric.trackers.get(metric.fullname)._metrics
+        assert metric.tracker._labelnames == ("pipeline",)
 
     def test_counter_metric_increments_correctly(self):
         metric = CounterMetric(
@@ -80,10 +85,11 @@ def test_no_duplicated_counter_is_created(self):
         )
         metric2.init_tracker()
 
-        assert metric1.trackers == metric2.trackers
+        assert metric1.tracker._labelnames == metric2.tracker._labelnames
         metric1 += 1
+        metric2 += 1
         metric_output = generate_latest(self.custom_registry).decode("utf-8")
-        result = re.findall(r'.*logprep_bla_total\{pipeline="1"\} 1\.0.*', metric_output)
+        result = re.findall(r'.*logprep_bla_total\{pipeline="1"\} 2\.0.*', metric_output)
         assert len(result) == 1
 
     def test_no_duplicated_counter_is_created_2(self):
@@ -102,30 +108,27 @@ def test_no_duplicated_counter_is_created_2(self):
         )
         metric2.init_tracker()
 
-        assert metric1.trackers == metric2.trackers
+        assert metric1.tracker == metric2.tracker
         metric1 += 1
         metric_output = generate_latest(self.custom_registry).decode("utf-8")
         result = re.findall(r'.*logprep_bla_total\{pipeline="1"\} 1\.0.*', metric_output)
         assert len(result) == 1
         result = re.findall(r'.*logprep_bla_total\{pipeline="2"\} 0\.0.*', metric_output)
         assert len(result) == 1
 
-    def test_tracker_contains_only_own_metric_types(self):
+    def test_init_tracker_raises_on_try_to_overwrite_tracker_with_different_type(self):
         metric1 = CounterMetric(
-            name="bla_counter",
+            name="bla",
             description="empty description",
             labels={"pipeline": "1"},
             registry=self.custom_registry,
         )
         metric1.init_tracker()
         metric2 = HistogramMetric(
-            name="bla_histogram",
+            name="bla",
             description="empty description",
             labels={"pipeline": "2"},
             registry=self.custom_registry,
         )
-        metric2.init_tracker()
-        assert len(metric1.trackers) == 1
-        assert len(metric2.trackers) == 1
-        assert isinstance(metric1.trackers.get(metric1.fullname), Counter)
-        assert isinstance(metric2.trackers.get(metric2.fullname), Histogram)
+        with pytest.raises(ValueError, match="already exists with different type"):
+            metric2.init_tracker()