Add tests for async functions

Signed-off-by: Yury Pliner <[email protected]>

Author: Pliner

tests/test_core.py

@@ -1,9 +1,9 @@
+import asyncio
+import aiounittest
 from concurrent.futures import ThreadPoolExecutor
 import time
 import unittest
 
-import pytest
-
 from prometheus_client.core import (
     CollectorRegistry, Counter, CounterMetricFamily, Enum, Gauge,
     GaugeHistogramMetricFamily, GaugeMetricFamily, Histogram,
@@ -28,7 +28,7 @@ def assert_not_observable(fn, *args, **kwargs):
     assert False, "Did not raise a 'missing label values' exception"
 
 
-class TestCounter(unittest.TestCase):
+class TestCounter(aiounittest.AsyncTestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
         self.counter = Counter('c_total', 'help', registry=self.registry)
@@ -56,30 +56,42 @@ def f(r):
 
         self.assertEqual((["r"], None, None, None), getargspec(f))
 
-        try:
+        with self.assertRaises(TypeError):
             f(False)
-        except TypeError:
-            pass
         self.assertEqual(0, self.registry.get_sample_value('c_total'))
 
-        try:
+        with self.assertRaises(ValueError):
             f(True)
-        except ValueError:
-            pass
+        self.assertEqual(1, self.registry.get_sample_value('c_total'))
+
+    async def test_async_function_decorator(self):
+        @self.counter.count_exceptions(ValueError)
+        async def f(r):
+            if r:
+                raise ValueError
+            else:
+                raise TypeError
+
+        self.assertEqual((["r"], None, None, None), getargspec(f))
+
+        with self.assertRaises(TypeError):
+            await f(False)
+
+        self.assertEqual(0, self.registry.get_sample_value('c_total'))
+
+        with self.assertRaises(ValueError):
+            await f(True)
+
         self.assertEqual(1, self.registry.get_sample_value('c_total'))
 
     def test_block_decorator(self):
         with self.counter.count_exceptions():
             pass
         self.assertEqual(0, self.registry.get_sample_value('c_total'))
 
-        raised = False
-        try:
+        with self.assertRaises(ValueError):
             with self.counter.count_exceptions():
                 raise ValueError
-        except:
-            raised = True
-        self.assertTrue(raised)
         self.assertEqual(1, self.registry.get_sample_value('c_total'))
 
     def test_count_exceptions_not_observable(self):
@@ -115,7 +127,7 @@ def test_exemplar_too_long(self):
         })
 
 
-class TestGauge(unittest.TestCase):
+class TestGauge(aiounittest.AsyncTestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
         self.gauge = Gauge('g', 'help', registry=self.registry)
@@ -160,6 +172,18 @@ def f():
         f()
         self.assertEqual(0, self.registry.get_sample_value('g'))
 
+    async def test_inprogress_async_function_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+
+        @self.gauge.track_inprogress()
+        async def f():
+            self.assertEqual(1, self.registry.get_sample_value('g'))
+
+        self.assertEqual(([], None, None, None), getargspec(f))
+
+        await f()
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+
     def test_inprogress_block_decorator(self):
         self.assertEqual(0, self.registry.get_sample_value('g'))
         with self.gauge.track_inprogress():
@@ -185,12 +209,24 @@ def test_time_function_decorator(self):
 
         @self.gauge.time()
         def f():
-            time.sleep(.001)
+            time.sleep(.05)
 
         self.assertEqual(([], None, None, None), getargspec(f))
 
         f()
-        self.assertNotEqual(0, self.registry.get_sample_value('g'))
+        self.assertTrue(0.05 <= self.registry.get_sample_value('g') <= 0.1)
+
+    async def test_time_async_function_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('g'))
+
+        @self.gauge.time()
+        async def f():
+            await asyncio.sleep(.05)
+
+        self.assertEqual(([], None, None, None), getargspec(f))
+
+        await f()
+        self.assertTrue(0.05 <= self.registry.get_sample_value('g') <= 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('g'))
@@ -239,7 +275,7 @@ def manager():
         assert_not_observable(manager)
 
 
-class TestSummary(unittest.TestCase):
+class TestSummary(aiounittest.AsyncTestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
         self.summary = Summary('s', 'help', registry=self.registry)
@@ -264,12 +300,26 @@ def test_function_decorator(self):
 
         @self.summary.time()
         def f():
-            pass
+            time.sleep(.05)
 
         self.assertEqual(([], None, None, None), getargspec(f))
 
         f()
         self.assertEqual(1, self.registry.get_sample_value('s_count'))
+        self.assertTrue(.05 < self.registry.get_sample_value('s_sum') < 0.1)
+
+    async def test_async_function_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('s_count'))
+
+        @self.summary.time()
+        async def f():
+            await asyncio.sleep(.05)
+
+        self.assertEqual(([], None, None, None), getargspec(f))
+
+        await f()
+        self.assertEqual(1, self.registry.get_sample_value('s_count'))
+        self.assertTrue(.05 < self.registry.get_sample_value('s_sum') < 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('s_count'))
@@ -343,7 +393,7 @@ def manager():
         assert_not_observable(manager)
 
 
-class TestHistogram(unittest.TestCase):
+class TestHistogram(aiounittest.AsyncTestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
         self.histogram = Histogram('h', 'help', registry=self.registry)
@@ -417,13 +467,29 @@ def test_function_decorator(self):
 
         @self.histogram.time()
         def f():
-            pass
+            time.sleep(.05)
 
         self.assertEqual(([], None, None, None), getargspec(f))
 
         f()
         self.assertEqual(1, self.registry.get_sample_value('h_count'))
         self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
+        self.assertTrue(.05 < self.registry.get_sample_value('h_sum') < 0.1)
+
+    async def test_async_function_decorator(self):
+        self.assertEqual(0, self.registry.get_sample_value('h_count'))
+        self.assertEqual(0, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
+
+        @self.histogram.time()
+        async def f():
+            await asyncio.sleep(.05)
+
+        self.assertEqual(([], None, None, None), getargspec(f))
+
+        await f()
+        self.assertEqual(1, self.registry.get_sample_value('h_count'))
+        self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
+        self.assertTrue(.05 < self.registry.get_sample_value('h_sum') < 0.1)
 
     def test_function_decorator_multithread(self):
         self.assertEqual(0, self.registry.get_sample_value('h_count'))
@@ -527,7 +593,7 @@ def test_labels(self):
         self.assertRaises(ValueError, self.labels.state, 'a')
 
     def test_overlapping_labels(self):
-        with pytest.raises(ValueError):
+        with self.assertRaises(ValueError):
             Enum('e', 'help', registry=None, labelnames=['e'])
 
 
@@ -568,7 +634,7 @@ def test_incorrect_label_count_raises(self):
         self.assertRaises(ValueError, self.counter.remove, 'a', 'b')
 
     def test_labels_on_labels(self):
-        with pytest.raises(ValueError):
+        with self.assertRaises(ValueError):
             self.counter.labels('a').labels('b')
 
     def test_labels_coerced_to_string(self):

tox.ini

@@ -6,6 +6,7 @@ envlist = coverage-clean,py3.6,py3.7,py3.8,py3.9,py3.10,pypy3.7,py3.9-nooptional
 deps =
     coverage
     pytest
+    aiounittest
     attrs
 
 [testenv]