mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-07-31 06:03:21 +08:00
304 lines
9.1 KiB
Python
304 lines
9.1 KiB
Python
# stdib
|
|
import time
|
|
import unittest
|
|
|
|
# 3p
|
|
import mongoengine
|
|
import pymongo
|
|
|
|
# project
|
|
from ddtrace import Pin
|
|
from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
|
|
from ddtrace.contrib.mongoengine.patch import patch, unpatch
|
|
from ddtrace.ext import mongo as mongox
|
|
|
|
# testing
|
|
from tests.opentracer.utils import init_tracer
|
|
from ..config import MONGO_CONFIG
|
|
from ...base import override_config
|
|
from ...test_tracer import get_dummy_tracer
|
|
|
|
|
|
class Artist(mongoengine.Document):
|
|
first_name = mongoengine.StringField(max_length=50)
|
|
last_name = mongoengine.StringField(max_length=50)
|
|
|
|
|
|
class MongoEngineCore(object):
|
|
|
|
# Define the service at the class level, so that each test suite can use a different service
|
|
# and therefore catch any sneaky badly-unpatched stuff.
|
|
TEST_SERVICE = 'deadbeef'
|
|
|
|
def get_tracer_and_connect(self):
|
|
# implement me
|
|
pass
|
|
|
|
def test_insert_update_delete_query(self):
|
|
tracer = self.get_tracer_and_connect()
|
|
|
|
start = time.time()
|
|
Artist.drop_collection()
|
|
end = time.time()
|
|
|
|
# ensure we get a drop collection span
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == 'drop artist'
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
start = end
|
|
joni = Artist()
|
|
joni.first_name = 'Joni'
|
|
joni.last_name = 'Mitchell'
|
|
joni.save()
|
|
end = time.time()
|
|
|
|
# ensure we get an insert span
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == 'insert artist'
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
# ensure full scans work
|
|
start = time.time()
|
|
artists = [a for a in Artist.objects]
|
|
end = time.time()
|
|
assert len(artists) == 1
|
|
assert artists[0].first_name == 'Joni'
|
|
assert artists[0].last_name == 'Mitchell'
|
|
|
|
# query names should be used in pymongo>3.1
|
|
name = 'find' if pymongo.version_tuple >= (3, 1, 0) else 'query'
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == '{} artist'.format(name)
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
# ensure filtered queries work
|
|
start = time.time()
|
|
artists = [a for a in Artist.objects(first_name='Joni')]
|
|
end = time.time()
|
|
assert len(artists) == 1
|
|
joni = artists[0]
|
|
assert artists[0].first_name == 'Joni'
|
|
assert artists[0].last_name == 'Mitchell'
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == '{} artist {{"first_name": "?"}}'.format(name)
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
# ensure updates work
|
|
start = time.time()
|
|
joni.last_name = 'From Saskatoon'
|
|
joni.save()
|
|
end = time.time()
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == 'update artist {"_id": "?"}'
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
# ensure deletes
|
|
start = time.time()
|
|
joni.delete()
|
|
end = time.time()
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
span = spans[0]
|
|
assert span.resource == 'delete artist {"_id": "?"}'
|
|
assert span.span_type == 'mongodb'
|
|
assert span.service == self.TEST_SERVICE
|
|
_assert_timing(span, start, end)
|
|
|
|
def test_opentracing(self):
|
|
"""Ensure the opentracer works with mongoengine."""
|
|
tracer = self.get_tracer_and_connect()
|
|
ot_tracer = init_tracer('my_svc', tracer)
|
|
|
|
with ot_tracer.start_active_span('ot_span'):
|
|
start = time.time()
|
|
Artist.drop_collection()
|
|
end = time.time()
|
|
|
|
# ensure we get a drop collection span
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 2
|
|
ot_span, dd_span = spans
|
|
|
|
# confirm the parenting
|
|
assert ot_span.parent_id is None
|
|
assert dd_span.parent_id == ot_span.span_id
|
|
|
|
assert ot_span.name == 'ot_span'
|
|
assert ot_span.service == 'my_svc'
|
|
|
|
assert dd_span.resource == 'drop artist'
|
|
assert dd_span.span_type == 'mongodb'
|
|
assert dd_span.service == self.TEST_SERVICE
|
|
_assert_timing(dd_span, start, end)
|
|
|
|
def test_analytics_default(self):
|
|
tracer = self.get_tracer_and_connect()
|
|
Artist.drop_collection()
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
assert spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY) is None
|
|
|
|
def test_analytics_with_rate(self):
|
|
with override_config(
|
|
'pymongo',
|
|
dict(analytics_enabled=True, analytics_sample_rate=0.5)
|
|
):
|
|
tracer = self.get_tracer_and_connect()
|
|
Artist.drop_collection()
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
assert spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY) == 0.5
|
|
|
|
def test_analytics_without_rate(self):
|
|
with override_config(
|
|
'pymongo',
|
|
dict(analytics_enabled=True)
|
|
):
|
|
tracer = self.get_tracer_and_connect()
|
|
Artist.drop_collection()
|
|
|
|
spans = tracer.writer.pop()
|
|
assert len(spans) == 1
|
|
assert spans[0].get_metric(ANALYTICS_SAMPLE_RATE_KEY) == 1.0
|
|
|
|
|
|
class TestMongoEnginePatchConnectDefault(unittest.TestCase, MongoEngineCore):
|
|
"""Test suite with a global Pin for the connect function with the default configuration"""
|
|
|
|
TEST_SERVICE = mongox.SERVICE
|
|
|
|
def setUp(self):
|
|
patch()
|
|
|
|
def tearDown(self):
|
|
unpatch()
|
|
# Disconnect and remove the client
|
|
mongoengine.connection.disconnect()
|
|
|
|
def get_tracer_and_connect(self):
|
|
tracer = get_dummy_tracer()
|
|
Pin.get_from(mongoengine.connect).clone(
|
|
tracer=tracer).onto(mongoengine.connect)
|
|
mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
|
|
return tracer
|
|
|
|
|
|
class TestMongoEnginePatchConnect(TestMongoEnginePatchConnectDefault):
|
|
"""Test suite with a global Pin for the connect function with custom service"""
|
|
|
|
TEST_SERVICE = 'test-mongo-patch-connect'
|
|
|
|
def get_tracer_and_connect(self):
|
|
tracer = TestMongoEnginePatchConnectDefault.get_tracer_and_connect(self)
|
|
Pin(service=self.TEST_SERVICE, tracer=tracer).onto(mongoengine.connect)
|
|
mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
|
|
return tracer
|
|
|
|
|
|
class TestMongoEnginePatchClientDefault(unittest.TestCase, MongoEngineCore):
|
|
"""Test suite with a Pin local to a specific client with default configuration"""
|
|
|
|
TEST_SERVICE = mongox.SERVICE
|
|
|
|
def setUp(self):
|
|
patch()
|
|
|
|
def tearDown(self):
|
|
unpatch()
|
|
# Disconnect and remove the client
|
|
mongoengine.connection.disconnect()
|
|
|
|
def get_tracer_and_connect(self):
|
|
tracer = get_dummy_tracer()
|
|
client = mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
Pin.get_from(client).clone(tracer=tracer).onto(client)
|
|
|
|
return tracer
|
|
|
|
|
|
class TestMongoEnginePatchClient(TestMongoEnginePatchClientDefault):
|
|
"""Test suite with a Pin local to a specific client with custom service"""
|
|
|
|
TEST_SERVICE = 'test-mongo-patch-client'
|
|
|
|
def get_tracer_and_connect(self):
|
|
tracer = get_dummy_tracer()
|
|
# Set a connect-level service, to check that we properly override it
|
|
Pin(service='not-%s' % self.TEST_SERVICE).onto(mongoengine.connect)
|
|
client = mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)
|
|
|
|
return tracer
|
|
|
|
def test_patch_unpatch(self):
|
|
tracer = get_dummy_tracer()
|
|
|
|
# Test patch idempotence
|
|
patch()
|
|
patch()
|
|
|
|
client = mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
Pin.get_from(client).clone(tracer=tracer).onto(client)
|
|
|
|
Artist.drop_collection()
|
|
spans = tracer.writer.pop()
|
|
assert spans, spans
|
|
assert len(spans) == 1
|
|
|
|
mongoengine.connection.disconnect()
|
|
tracer.writer.pop()
|
|
|
|
# Test unpatch
|
|
unpatch()
|
|
|
|
mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
|
|
Artist.drop_collection()
|
|
spans = tracer.writer.pop()
|
|
assert not spans, spans
|
|
|
|
# Test patch again
|
|
patch()
|
|
|
|
client = mongoengine.connect(port=MONGO_CONFIG['port'])
|
|
Pin.get_from(client).clone(tracer=tracer).onto(client)
|
|
|
|
Artist.drop_collection()
|
|
spans = tracer.writer.pop()
|
|
assert spans, spans
|
|
assert len(spans) == 1
|
|
|
|
|
|
def _assert_timing(span, start, end):
|
|
assert start < span.start < end
|
|
assert span.duration < end - start
|