Fix async redis clients tracing (#1830)

* Fix async redis clients tracing

* Update changelog

* Add functional integration tests and fix linting issues

---------

Co-authored-by: Shalev Roda <65566801+shalevr@users.noreply.github.com>
This commit is contained in:
Vivanov98
2023-06-25 13:03:54 +01:00
committed by GitHub
parent e70437a36e
commit cd6b024327
4 changed files with 270 additions and 34 deletions

View File

@ -11,9 +11,11 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
from unittest import mock
import redis
import redis.asyncio
from opentelemetry import trace
from opentelemetry.instrumentation.redis import RedisInstrumentor
@ -21,6 +23,24 @@ from opentelemetry.test.test_base import TestBase
from opentelemetry.trace import SpanKind
class AsyncMock:
"""A sufficient async mock implementation.
Python 3.7 doesn't have an inbuilt async mock class, so this is used.
"""
def __init__(self):
self.mock = mock.Mock()
async def __call__(self, *args, **kwargs):
future = asyncio.Future()
future.set_result("random")
return future
def __getattr__(self, item):
return AsyncMock()
class TestRedis(TestBase):
def setUp(self):
super().setUp()
@ -87,6 +107,35 @@ class TestRedis(TestBase):
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 1)
def test_instrument_uninstrument_async_client_command(self):
redis_client = redis.asyncio.Redis()
with mock.patch.object(redis_client, "connection", AsyncMock()):
asyncio.run(redis_client.get("key"))
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 1)
self.memory_exporter.clear()
# Test uninstrument
RedisInstrumentor().uninstrument()
with mock.patch.object(redis_client, "connection", AsyncMock()):
asyncio.run(redis_client.get("key"))
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 0)
self.memory_exporter.clear()
# Test instrument again
RedisInstrumentor().instrument()
with mock.patch.object(redis_client, "connection", AsyncMock()):
asyncio.run(redis_client.get("key"))
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 1)
def test_response_hook(self):
redis_client = redis.Redis()
connection = redis.connection.Connection()