Files
2020-04-08 10:39:44 -07:00

88 lines
2.5 KiB
Python

from functools import wraps
# 3rd party
from django.apps import apps
from django.test import TestCase
# project
from ddtrace.tracer import Tracer
from ddtrace.contrib.django.conf import settings
from ddtrace.contrib.django.db import patch_db, unpatch_db
from ddtrace.contrib.django.cache import unpatch_cache
from ddtrace.contrib.django.templates import unpatch_template
from ddtrace.contrib.django.middleware import remove_exception_middleware, remove_trace_middleware
# testing
from ...base import BaseTestCase
from ...test_tracer import DummyWriter
# testing tracer
tracer = Tracer()
tracer.writer = DummyWriter()
class DjangoTraceTestCase(BaseTestCase, TestCase):
"""
Base class that provides an internal tracer according to given
Datadog settings. This class ensures that the tracer spans are
properly reset after each run. The tracer is available in
the ``self.tracer`` attribute.
"""
def setUp(self):
# assign the default tracer
self.tracer = settings.TRACER
# empty the tracer spans from previous operations
# such as database creation queries
self.tracer.writer.spans = []
self.tracer.writer.pop_traces()
# gets unpatched for some tests
patch_db(self.tracer)
def tearDown(self):
# empty the tracer spans from test operations
self.tracer.writer.spans = []
self.tracer.writer.pop_traces()
class override_ddtrace_settings(object):
def __init__(self, *args, **kwargs):
self.items = list(kwargs.items())
def unpatch_all(self):
unpatch_cache()
unpatch_db()
unpatch_template()
remove_trace_middleware()
remove_exception_middleware()
def __enter__(self):
self.enable()
def __exit__(self, exc_type, exc_value, traceback):
self.disable()
def enable(self):
self.backup = {}
for name, value in self.items:
self.backup[name] = getattr(settings, name)
setattr(settings, name, value)
self.unpatch_all()
app = apps.get_app_config('datadog_django')
app.ready()
def disable(self):
for name, value in self.items:
setattr(settings, name, self.backup[name])
self.unpatch_all()
remove_exception_middleware()
app = apps.get_app_config('datadog_django')
app.ready()
def __call__(self, func):
@wraps(func)
def inner(*args, **kwargs):
with(self):
return func(*args, **kwargs)
return inner