Rewrite the KeyRing (#10035)

1 files changed, 16 insertions, 21 deletions

diff --git a/tests/util/test_batching_queue.py b/tests/util/test_batching_queue.py
index edf29e5b96..07be57d72c 100644
--- a/tests/util/test_batching_queue.py
+++ b/tests/util/test_batching_queue.py
@@ -45,37 +45,32 @@ class BatchingQueueTestCase(TestCase):
         self._pending_calls.append((values, d))
         return await make_deferred_yieldable(d)
 
+    def _get_sample_with_name(self, metric, name) -> int:
+        """For a prometheus metric get the value of the sample that has a
+        matching "name" label.
+        """
+        for sample in metric.collect()[0].samples:
+            if sample.labels.get("name") == name:
+                return sample.value
+
+        self.fail("Found no matching sample")
+
     def _assert_metrics(self, queued, keys, in_flight):
         """Assert that the metrics are correct"""
 
-        self.assertEqual(len(number_queued.collect()), 1)
-        self.assertEqual(len(number_queued.collect()[0].samples), 1)
+        sample = self._get_sample_with_name(number_queued, self.queue._name)
         self.assertEqual(
-            number_queued.collect()[0].samples[0].labels,
-            {"name": self.queue._name},
-        )
-        self.assertEqual(
-            number_queued.collect()[0].samples[0].value,
+            sample,
             queued,
             "number_queued",
         )
 
-        self.assertEqual(len(number_of_keys.collect()), 1)
-        self.assertEqual(len(number_of_keys.collect()[0].samples), 1)
-        self.assertEqual(
-            number_queued.collect()[0].samples[0].labels, {"name": self.queue._name}
-        )
-        self.assertEqual(
-            number_of_keys.collect()[0].samples[0].value, keys, "number_of_keys"
-        )
+        sample = self._get_sample_with_name(number_of_keys, self.queue._name)
+        self.assertEqual(sample, keys, "number_of_keys")
 
-        self.assertEqual(len(number_in_flight.collect()), 1)
-        self.assertEqual(len(number_in_flight.collect()[0].samples), 1)
-        self.assertEqual(
-            number_queued.collect()[0].samples[0].labels, {"name": self.queue._name}
-        )
+        sample = self._get_sample_with_name(number_in_flight, self.queue._name)
         self.assertEqual(
-            number_in_flight.collect()[0].samples[0].value,
+            sample,
             in_flight,
             "number_in_flight",
         )