fix enum metrics

tangostationcontrol/tangostationcontrol/devices/calibration.py

@@ -57,8 +57,9 @@ class Calibration(LOFARDevice):
         self.calibration_count_metric = AttributeMetric(
             "calibration_count",
             "Number of times calibration has been triggered for each AntennaField device",
-            device_labels(self) + ["antennafield"],
+            device_labels(self),
             Counter,
+            dynamic_labels=["antennafield"],
         )
 
     def _calibrate_antenna_field(self, device):
@@ -78,7 +79,7 @@ class Calibration(LOFARDevice):
         # get device member in its original casing
         antenna_field_name = device.get_name().split("/")[2]
 
-        self.calibration_count_metric.get_metric(antenna_field_name).inc()
+        self.calibration_count_metric.get_metric([antenna_field_name]).inc()
 
     @log_exceptions()
     def _antennafield_changed_event(self, event):

tangostationcontrol/tangostationcontrol/metrics/_decorators.py

@@ -38,22 +38,17 @@ class VersionMetric(AttributeMetric):
 
 class StateMetric(AttributeMetric):
     def __init__(self, device: Device):
-        super().__init__("state", "State of the device.", device.metric_labels, Enum)
-
-        self.set_state(device.get_state())
-
-        wrap_method(device, device.set_state, self.set_state, post_execute=False)
-
-    def make_metric(self) -> Metric:
-        return Enum(
-            self.name,
-            self.description,
-            labelnames=self.label_keys(),
-            states=list(DevState.names),
+        super().__init__(
+            "state",
+            "State of the device.",
+            device.metric_labels,
+            Enum,
+            metric_class_init_kwargs={"states": list(DevState.names)},
         )
 
-    def set_state(self, state):
-        self.get_metric().state(state.name)
+        self.set_value(device.get_state())
+
+        wrap_method(device, device.set_state, self.set_value, post_execute=False)
 
 
 class AccessCountMetric(AttributeMetric):

tangostationcontrol/tangostationcontrol/metrics/_metrics.py

@@ -3,8 +3,10 @@ from tango import AttrWriteType
 from tango import CmdArgType
 from tango import Attribute
 from tango import DevFailed
-from prometheus_client import Metric, Gauge, Info
+from tango import DevState
+from prometheus_client import Metric, Gauge, Info, Enum
 from asyncio import iscoroutinefunction
+from enum import IntEnum
 from typing import List, Dict, Callable, Union
 import functools
 import logging
@@ -112,6 +114,8 @@ class AttributeMetric:
         description: str,
         static_labels: Dict[str, str],
         metric_class=Gauge,
+        metric_class_init_kwargs: Dict[str, object] | None = None,
+        dynamic_labels: List[str] | None = None,
     ):
         self.name = metric_name(name)
         self.description = description
@@ -120,6 +124,10 @@ class AttributeMetric:
         self.static_label_keys = list(static_labels.keys())
         self.static_label_values = list(static_labels.values())
 
+        self.dynamic_label_keys = dynamic_labels or []
+
+        self.metric_class_init_kwargs = metric_class_init_kwargs or {}
+
         if self.name not in METRICS:
             METRICS[self.name] = self.make_metric()
 
@@ -131,22 +139,30 @@ class AttributeMetric:
     def label_keys(self) -> List[str]:
         """Return the list of labels that we will use."""
 
-        return self.static_label_keys
+        return self.static_label_keys + self.dynamic_label_keys
 
     def make_metric(self) -> Metric:
         """Construct a metric that collects samples for this attribute."""
         return self.metric_class(
-            self.name, self.description, labelnames=self.label_keys()
+            self.name,
+            self.description,
+            labelnames=self.label_keys(),
+            **self.metric_class_init_kwargs,
         )
 
-    def get_metric(self, extra_labels: List = None) -> Metric:
+    def get_metric(self, dynamic_label_values: List = None) -> Metric:
         """Return the metric that uses the default labels."""
-        return self.metric.labels(*self.static_label_values, *(extra_labels or []))
+        return self.metric.labels(
+            *self.static_label_values, *(dynamic_label_values or [])
+        )
 
     def set_value(self, value: object):
         """A new value for the attribute is known. Feed it to the metric."""
 
         # set it, this class will take care of the default labels
+        if self.metric_class == Enum:
+            self._enum_value(value, self.static_label_values)
+        else:
             self._set_value(value, self.static_label_values)
 
     def _set_value(self, value: object, labels: List[str]):
@@ -157,6 +173,12 @@ class AttributeMetric:
         assert self.metric_class == Info
         self.metric.labels(*labels).info(value)
 
+    def _enum_value(self, value: str | IntEnum, labels: List[str]):
+        assert self.metric_class == Enum
+        self.metric.labels(*labels).state(
+            value.name if isinstance(value, (DevState, IntEnum)) else value
+        )
+
     def collect(self) -> List[Metric]:
         """Return all collected samples."""
         return self.metric.collect()
@@ -183,6 +205,19 @@ class ScalarAttributeMetric(AttributeMetric):
 
         if self.data_type == CmdArgType.DevString:
             super().__init__(attribute.get_name(), description, static_labels, Info)
+        elif self.data_type == CmdArgType.DevEnum:
+            # evil PyTango foo to obtain enum labels from class attribute
+            enum_labels = getattr(
+                device.__class__, attribute.get_name()
+            ).att_prop.enum_labels.split(",")
+
+            super().__init__(
+                attribute.get_name(),
+                description,
+                static_labels,
+                Enum,
+                metric_class_init_kwargs={"states": enum_labels},
+            )
         else:
             super().__init__(attribute.get_name(), description, static_labels)
 

tangostationcontrol/test/metrics/test_metrics.py

@@ -10,6 +10,7 @@ from tango.server import (
 from tango.test_context import DeviceTestContext
 from prometheus_client import generate_latest
 from prometheus_client.registry import REGISTRY
+from enum import IntEnum
 from typing import Dict
 import asyncio
 import numpy
@@ -107,6 +108,10 @@ class TestMetrics(base.TestCase):
     def test_scalar_attribute_metric(self):
         """Test ScalarAttributeMetric"""
 
+        class MyEnum(IntEnum):
+            ZERO = 0
+            ONE = 1
+
         class test_device(Device):
             float_attr = attribute(
                 doc="docstr",
@@ -120,10 +125,17 @@ class TestMetrics(base.TestCase):
                 fget=lambda obj: "foo",
             )
 
+            enum_attr = attribute(
+                doc="docstr",
+                dtype=MyEnum,
+                fget=lambda obj: MyEnum.ONE,
+            )
+
             def init_device(self):
                 # create an attribute metric and assign a value
                 self.float_metric = ScalarAttributeMetric(self, self.float_attr)
                 self.str_metric = ScalarAttributeMetric(self, self.str_attr)
+                self.enum_metric = ScalarAttributeMetric(self, self.enum_attr)
 
             @command()
             def test(device):
@@ -166,10 +178,44 @@ class TestMetrics(base.TestCase):
                     metric.samples[0].labels,
                 )
 
+                # check collected metrics (enum_attr)
+                metric = device.enum_metric.metric.collect()[0]
+
+                self.assertEqual("ds_enum_attr", metric.name)
+                self.assertEqual("docstr", metric.documentation)
+
+                # check labels as the DeviceTestContext would result in
+                self.assertDictEqual(
+                    {
+                        "domain": "test",
+                        "family": "nodb",
+                        "member": "test_device",
+                        "device_class": "test_device",
+                        "access": "r",
+                        "ds_enum_attr": "ZERO",
+                    },
+                    metric.samples[0].labels,
+                )
+                self.assertEqual(1, metric.samples[0].value)
+
+                self.assertDictEqual(
+                    {
+                        "domain": "test",
+                        "family": "nodb",
+                        "member": "test_device",
+                        "device_class": "test_device",
+                        "access": "r",
+                        "ds_enum_attr": "ONE",
+                    },
+                    metric.samples[1].labels,
+                )
+                self.assertEqual(0, metric.samples[1].value)
+
         with DeviceTestContext(test_device, process=False) as proxy:
             # access the attribute to trigger value propagation to metric
             _ = proxy.float_attr
             _ = proxy.str_attr
+            _ = proxy.enum_attr
 
             proxy.test()