Add back the assert_span_count method

Author: OlegZv

instrumentation/opentelemetry-instrumentation-redis/tests/test_redis.py

@@ -35,6 +35,15 @@
 
 
 class TestRedis(TestBase):
+ def assert_span_count(self, count: int):
+ """
+ Assert that the memory exporter has the expected number of spans.
+ Returns the spans list if assertion passes
+ """
+ spans = self.memory_exporter.get_finished_spans()
+ self.assertEqual(len(spans), count)
+ return spans
+
  def setUp(self):
  super().setUp()
  RedisInstrumentor().instrument(tracer_provider=self.tracer_provider)
@@ -387,6 +396,15 @@ def redis_operations():
 
 
 class TestRedisAsync(TestBase, IsolatedAsyncioTestCase):
+ def assert_span_count(self, count: int):
+ """
+ Assert that the memory exporter has the expected number of spans.
+ Returns the spans list if assertion passes
+ """
+ spans = self.memory_exporter.get_finished_spans()
+ self.assertEqual(len(spans), count)
+ return spans
+
  def setUp(self):
  super().setUp()
  self.instrumentor = RedisInstrumentor()
@@ -530,6 +548,15 @@ def response_hook(span, conn, args):
 
 
 class TestRedisInstance(TestBase):
+ def assert_span_count(self, count: int):
+ """
+ Assert that the memory exporter has the expected number of spans.
+ Returns the spans list if assertion passes
+ """
+ spans = self.memory_exporter.get_finished_spans()
+ self.assertEqual(len(spans), count)
+ return spans
+
  def setUp(self):
  super().setUp()
  self.client = fakeredis.FakeStrictRedis()