From a2f0255c7654da042f77b2cb49b6e32428d62967 Mon Sep 17 00:00:00 2001 From: Sean Anderson Date: Fri, 17 Nov 2023 20:06:57 -0500 Subject: [PATCH] Fix histograms not being unpickled correctly _Histogram._setstate does not pass heap to Struct._setstate, which caused it to be initialized to None. This results in exemplars breaking, since they access heap. Fix this, and update the unpickling test to be more functional. Fixes: 78f0ce7 ("types: Rework state-setting functions") Signed-off-by: Sean Anderson --- mpmetrics/metrics.py | 2 +- mpmetrics/types.py | 1 + tests/metric_test.py | 3 ++- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/mpmetrics/metrics.py b/mpmetrics/metrics.py index 4c58d3c..a4b599b 100644 --- a/mpmetrics/metrics.py +++ b/mpmetrics/metrics.py @@ -386,7 +386,7 @@ def __init__(self, mem, thresholds, heap, **kwargs): self._created.value = time.time() def _setstate(self, mem, heap): - Struct._setstate(self, mem) + Struct._setstate(self, mem, heap) self.thresholds = tuple(threshold.value for threshold in self._thresholds) def observe(self, amount, exemplar=None): diff --git a/mpmetrics/types.py b/mpmetrics/types.py index 47d8a0b..0fcfc4b 100644 --- a/mpmetrics/types.py +++ b/mpmetrics/types.py @@ -150,6 +150,7 @@ def __init__(self, mem, heap): self._heap = heap def _setstate(self, mem, heap): + assert heap is not None super()._setstate(mem, heap) self._heap = heap diff --git a/tests/metric_test.py b/tests/metric_test.py index 79e2321..cb9ac2f 100644 --- a/tests/metric_test.py +++ b/tests/metric_test.py @@ -349,7 +349,8 @@ def test_time(registry, cls, name): @pytest.mark.parametrize('cls', (Counter, Gauge, Summary, Histogram)) def test_pickle(registry, cls): metric = cls('name', 'help', labelnames=('l'), registry=registry) - pickle.loads(pickle.dumps(metric.labels('x'))) + metric.labels('x') + assert list(metric.collect()) == list(pickle.loads(pickle.dumps(metric)).collect()) class TestEnum: @pytest.fixture