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 <[email protected]>
Forty-Bot · Nov 18, 2023 · a2f0255 · a2f0255
1 parent 3a68e8e
commit a2f0255
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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