opentelemetry-python-contrib/reference/tests/test_sampler.py

from __future__ import division
import contextlib
import mock
import re
import unittest

import pytest

from ddtrace.compat import iteritems
from ddtrace.constants import SAMPLING_PRIORITY_KEY, SAMPLE_RATE_METRIC_KEY
from ddtrace.constants import SAMPLING_AGENT_DECISION, SAMPLING_RULE_DECISION, SAMPLING_LIMIT_DECISION
from ddtrace.ext.priority import AUTO_KEEP, AUTO_REJECT
from ddtrace.internal.rate_limiter import RateLimiter
from ddtrace.sampler import DatadogSampler, SamplingRule
from ddtrace.sampler import RateSampler, AllSampler, RateByServiceSampler
from ddtrace.span import Span

from .utils import override_env
from .test_tracer import get_dummy_tracer


@pytest.fixture
def dummy_tracer():
    return get_dummy_tracer()


def assert_sampling_decision_tags(span, agent=None, limit=None, rule=None):
    assert span.get_metric(SAMPLING_AGENT_DECISION) == agent
    assert span.get_metric(SAMPLING_LIMIT_DECISION) == limit
    assert span.get_metric(SAMPLING_RULE_DECISION) == rule


def create_span(tracer=None, name='test.span', meta=None, *args, **kwargs):
    tracer = tracer or get_dummy_tracer()
    if 'context' not in kwargs:
        kwargs['context'] = tracer.get_call_context()
    span = Span(tracer=tracer, name=name, *args, **kwargs)
    if meta:
        span.set_tags(meta)
    return span


class RateSamplerTest(unittest.TestCase):

    def test_set_sample_rate(self):
        sampler = RateSampler()
        assert sampler.sample_rate == 1.0

        for rate in [0.001, 0.01, 0.1, 0.25, 0.5, 0.75, 0.99999999, 1.0, 1]:
            sampler.set_sample_rate(rate)
            assert sampler.sample_rate == float(rate)

            sampler.set_sample_rate(str(rate))
            assert sampler.sample_rate == float(rate)

    def test_set_sample_rate_str(self):
        sampler = RateSampler()
        sampler.set_sample_rate('0.5')
        assert sampler.sample_rate == 0.5

    def test_sample_rate_deviation(self):
        for sample_rate in [0.1, 0.25, 0.5, 1]:
            tracer = get_dummy_tracer()
            writer = tracer.writer

            tracer.sampler = RateSampler(sample_rate)

            iterations = int(1e4 / sample_rate)

            for i in range(iterations):
                span = tracer.trace(i)
                span.finish()

            samples = writer.pop()

            # We must have at least 1 sample, check that it has its sample rate properly assigned
            assert samples[0].get_metric(SAMPLE_RATE_METRIC_KEY) == sample_rate

            # Less than 5% deviation when 'enough' iterations (arbitrary, just check if it converges)
            deviation = abs(len(samples) - (iterations * sample_rate)) / (iterations * sample_rate)
            assert deviation < 0.05, 'Deviation too high %f with sample_rate %f' % (deviation, sample_rate)

    def test_deterministic_behavior(self):
        """ Test that for a given trace ID, the result is always the same """
        tracer = get_dummy_tracer()
        writer = tracer.writer

        tracer.sampler = RateSampler(0.5)

        for i in range(10):
            span = tracer.trace(i)
            span.finish()

            samples = writer.pop()
            assert len(samples) <= 1, 'there should be 0 or 1 spans'
            sampled = (1 == len(samples))
            for j in range(10):
                other_span = Span(tracer, i, trace_id=span.trace_id)
                assert (
                    sampled == tracer.sampler.sample(other_span)
                ), 'sampling should give the same result for a given trace_id'


class RateByServiceSamplerTest(unittest.TestCase):
    def test_default_key(self):
        assert (
            'service:,env:' == RateByServiceSampler._default_key
        ), 'default key should correspond to no service and no env'

    def test_key(self):
        assert RateByServiceSampler._default_key == RateByServiceSampler._key()
        assert 'service:mcnulty,env:' == RateByServiceSampler._key(service='mcnulty')
        assert 'service:,env:test' == RateByServiceSampler._key(env='test')
        assert 'service:mcnulty,env:test' == RateByServiceSampler._key(service='mcnulty', env='test')
        assert 'service:mcnulty,env:test' == RateByServiceSampler._key('mcnulty', 'test')

    def test_sample_rate_deviation(self):
        for sample_rate in [0.1, 0.25, 0.5, 1]:
            tracer = get_dummy_tracer()
            writer = tracer.writer
            tracer.configure(sampler=AllSampler())
            # We need to set the writer because tracer.configure overrides it,
            # indeed, as we enable priority sampling, we must ensure the writer
            # is priority sampling aware and pass it a reference on the
            # priority sampler to send the feedback it gets from the agent
            assert writer != tracer.writer, 'writer should have been updated by configure'
            tracer.writer = writer
            tracer.priority_sampler.set_sample_rate(sample_rate)

            iterations = int(1e4 / sample_rate)

            for i in range(iterations):
                span = tracer.trace(i)
                span.finish()

            samples = writer.pop()
            samples_with_high_priority = 0
            for sample in samples:
                if sample.get_metric(SAMPLING_PRIORITY_KEY) is not None:
                    if sample.get_metric(SAMPLING_PRIORITY_KEY) > 0:
                        samples_with_high_priority += 1
                else:
                    assert (
                        0 == sample.get_metric(SAMPLING_PRIORITY_KEY)
                    ), 'when priority sampling is on, priority should be 0 when trace is to be dropped'
                assert_sampling_decision_tags(sample, agent=sample_rate)
            # We must have at least 1 sample, check that it has its sample rate properly assigned
            assert samples[0].get_metric(SAMPLE_RATE_METRIC_KEY) is None

            # Less than 5% deviation when 'enough' iterations (arbitrary, just check if it converges)
            deviation = abs(samples_with_high_priority - (iterations * sample_rate)) / (iterations * sample_rate)
            assert deviation < 0.05, 'Deviation too high %f with sample_rate %f' % (deviation, sample_rate)

    def test_update_rate_by_service_sample_rates(self):
        cases = [
            {
                'service:,env:': 1,
            },
            {
                'service:,env:': 1,
                'service:mcnulty,env:dev': 0.33,
                'service:postgres,env:dev': 0.7,
            },
            {
                'service:,env:': 1,
                'service:mcnulty,env:dev': 0.25,
                'service:postgres,env:dev': 0.5,
                'service:redis,env:prod': 0.75,
            },
        ]

        tracer = get_dummy_tracer()
        tracer.configure(sampler=AllSampler())
        priority_sampler = tracer.priority_sampler
        for case in cases:
            priority_sampler.update_rate_by_service_sample_rates(case)
            rates = {}
            for k, v in iteritems(priority_sampler._by_service_samplers):
                rates[k] = v.sample_rate
            assert case == rates, '%s != %s' % (case, rates)
        # It's important to also test in reverse mode for we want to make sure key deletion
        # works as well as key insertion (and doing this both ways ensures we trigger both cases)
        cases.reverse()
        for case in cases:
            priority_sampler.update_rate_by_service_sample_rates(case)
            rates = {}
            for k, v in iteritems(priority_sampler._by_service_samplers):
                rates[k] = v.sample_rate
            assert case == rates, '%s != %s' % (case, rates)


@pytest.mark.parametrize(
    'sample_rate,allowed',
    [
        # Min/max allowed values
        (0.0, True),
        (1.0, True),

        # Accepted boundaries
        (0.000001, True),
        (0.999999, True),

        # Outside the bounds
        (-0.000000001, False),
        (1.0000000001, False),
    ] + [
        # Try a bunch of decimal values between 0 and 1
        (1 / i, True) for i in range(1, 50)
    ] + [
        # Try a bunch of decimal values less than 0
        (-(1 / i), False) for i in range(1, 50)
    ] + [
        # Try a bunch of decimal values greater than 1
        (1 + (1 / i), False) for i in range(1, 50)
    ]
)
def test_sampling_rule_init_sample_rate(sample_rate, allowed):
    if allowed:
        rule = SamplingRule(sample_rate=sample_rate)
        assert rule.sample_rate == sample_rate
    else:
        with pytest.raises(ValueError):
            SamplingRule(sample_rate=sample_rate)


def test_sampling_rule_init_defaults():
    rule = SamplingRule(sample_rate=1.0)
    assert rule.sample_rate == 1.0
    assert rule.service == SamplingRule.NO_RULE
    assert rule.name == SamplingRule.NO_RULE


def test_sampling_rule_init():
    name_regex = re.compile(r'\.request$')

    def resource_check(resource):
        return 'healthcheck' in resource

    rule = SamplingRule(
        sample_rate=0.0,
        # Value
        service='my-service',
        # Regex
        name=name_regex,
    )

    assert rule.sample_rate == 0.0
    assert rule.service == 'my-service'
    assert rule.name == name_regex


@pytest.mark.parametrize(
    'span,rule,expected',
    [
        # DEV: Use sample_rate=1 to ensure SamplingRule._sample always returns True
        (create_span(name=name), SamplingRule(
            sample_rate=1, name=pattern), expected)
        for name, pattern, expected in [
            ('test.span', SamplingRule.NO_RULE, True),
            # DEV: `span.name` cannot be `None`
            ('test.span', None, False),
            ('test.span', 'test.span', True),
            ('test.span', 'test_span', False),
            ('test.span', re.compile(r'^test\.span$'), True),
            ('test_span', re.compile(r'^test.span$'), True),
            ('test.span', re.compile(r'^test_span$'), False),
            ('test.span', re.compile(r'test'), True),
            ('test.span', re.compile(r'test\.span|another\.span'), True),
            ('another.span', re.compile(r'test\.span|another\.span'), True),
            ('test.span', lambda name: 'span' in name, True),
            ('test.span', lambda name: 'span' not in name, False),
            ('test.span', lambda name: 1 / 0, False),
        ]
    ]
)
def test_sampling_rule_matches_name(span, rule, expected):
    assert rule.matches(span) is expected, '{} -> {} -> {}'.format(rule, span, expected)


@pytest.mark.parametrize(
    'span,rule,expected',
    [
        # DEV: Use sample_rate=1 to ensure SamplingRule._sample always returns True
        (create_span(service=service), SamplingRule(sample_rate=1, service=pattern), expected)
        for service, pattern, expected in [
            ('my-service', SamplingRule.NO_RULE, True),
            ('my-service', None, False),
            (None, None, True),
            (None, 'my-service', False),
            (None, re.compile(r'my-service'), False),
            (None, lambda service: 'service' in service, False),
            ('my-service', 'my-service', True),
            ('my-service', 'my_service', False),
            ('my-service', re.compile(r'^my-'), True),
            ('my_service', re.compile(r'^my[_-]'), True),
            ('my-service', re.compile(r'^my_'), False),
            ('my-service', re.compile(r'my-service'), True),
            ('my-service', re.compile(r'my'), True),
            ('my-service', re.compile(r'my-service|another-service'), True),
            ('another-service', re.compile(r'my-service|another-service'), True),
            ('my-service', lambda service: 'service' in service, True),
            ('my-service', lambda service: 'service' not in service, False),
            ('my-service', lambda service: 1 / 0, False),
        ]
    ]
)
def test_sampling_rule_matches_service(span, rule, expected):
    assert rule.matches(span) is expected, '{} -> {} -> {}'.format(rule, span, expected)


@pytest.mark.parametrize(
    'span,rule,expected',
    [
        # All match
        (
            create_span(
                name='test.span',
                service='my-service',
            ),
            SamplingRule(
                sample_rate=1,
                name='test.span',
                service=re.compile(r'^my-'),
            ),
            True,
        ),

        # All match,  but sample rate of 0%
        # DEV: We are checking if it is a match, not computing sampling rate, sample_rate=0 is not considered
        (
            create_span(
                name='test.span',
                service='my-service',
            ),
            SamplingRule(
                sample_rate=0,
                name='test.span',
                service=re.compile(r'^my-'),
            ),
            True,
        ),

        # Name doesn't match
        (
            create_span(
                name='test.span',
                service='my-service',
            ),
            SamplingRule(
                sample_rate=1,
                name='test_span',
                service=re.compile(r'^my-'),
            ),
            False,
        ),

        # Service doesn't match
        (
            create_span(
                name='test.span',
                service='my-service',
            ),
            SamplingRule(
                sample_rate=1,
                name='test.span',
                service=re.compile(r'^service-'),
            ),
            False,
        ),
    ],
)
def test_sampling_rule_matches(span, rule, expected):
    assert rule.matches(span) is expected, '{} -> {} -> {}'.format(rule, span, expected)


def test_sampling_rule_matches_exception():
    e = Exception('an error occurred')

    def pattern(prop):
        raise e

    rule = SamplingRule(sample_rate=1.0, name=pattern)
    span = create_span(name='test.span')

    with mock.patch('ddtrace.sampler.log') as mock_log:
        assert rule.matches(span) is False
        mock_log.warning.assert_called_once_with(
            '%r pattern %r failed with %r',
            rule,
            pattern,
            'test.span',
            exc_info=True,
        )


@pytest.mark.parametrize('sample_rate', [0.01, 0.1, 0.15, 0.25, 0.5, 0.75, 0.85, 0.9, 0.95, 0.991])
def test_sampling_rule_sample(sample_rate):
    tracer = get_dummy_tracer()
    rule = SamplingRule(sample_rate=sample_rate)

    iterations = int(1e4 / sample_rate)
    sampled = sum(
        rule.sample(Span(tracer=tracer, name=i))
        for i in range(iterations)
    )

    # Less than 5% deviation when 'enough' iterations (arbitrary, just check if it converges)
    deviation = abs(sampled - (iterations * sample_rate)) / (iterations * sample_rate)
    assert deviation < 0.05, (
        'Deviation {!r} too high with sample_rate {!r} for {} sampled'.format(deviation, sample_rate, sampled)
    )


def test_sampling_rule_sample_rate_1():
    tracer = get_dummy_tracer()
    rule = SamplingRule(sample_rate=1)

    iterations = int(1e4)
    assert all(
        rule.sample(Span(tracer=tracer, name=i))
        for i in range(iterations)
    )


def test_sampling_rule_sample_rate_0():
    tracer = get_dummy_tracer()
    rule = SamplingRule(sample_rate=0)

    iterations = int(1e4)
    assert sum(
        rule.sample(Span(tracer=tracer, name=i))
        for i in range(iterations)
    ) == 0


def test_datadog_sampler_init():
    # No args
    sampler = DatadogSampler()
    assert sampler.rules == []
    assert isinstance(sampler.limiter, RateLimiter)
    assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
    assert isinstance(sampler.default_sampler, RateByServiceSampler)

    # With rules
    rule = SamplingRule(sample_rate=1)
    sampler = DatadogSampler(rules=[rule])
    assert sampler.rules == [rule]
    assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
    assert isinstance(sampler.default_sampler, RateByServiceSampler)

    # With rate limit
    sampler = DatadogSampler(rate_limit=10)
    assert sampler.limiter.rate_limit == 10
    assert isinstance(sampler.default_sampler, RateByServiceSampler)

    # With default_sample_rate
    sampler = DatadogSampler(default_sample_rate=0.5)
    assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
    assert isinstance(sampler.default_sampler, SamplingRule)
    assert sampler.default_sampler.sample_rate == 0.5

    # From env variables
    with override_env(dict(DD_TRACE_SAMPLE_RATE='0.5', DD_TRACE_RATE_LIMIT='10')):
        sampler = DatadogSampler()
        assert sampler.limiter.rate_limit == 10
        assert isinstance(sampler.default_sampler, SamplingRule)
        assert sampler.default_sampler.sample_rate == 0.5

    # Invalid rules
    for val in (None, True, False, object(), 1, Exception()):
        with pytest.raises(TypeError):
            DatadogSampler(rules=[val])

    # Ensure rule order
    rule_1 = SamplingRule(sample_rate=1)
    rule_2 = SamplingRule(sample_rate=0.5, service='test')
    rule_3 = SamplingRule(sample_rate=0.25, name='flask.request')
    sampler = DatadogSampler(rules=[rule_1, rule_2, rule_3])
    assert sampler.rules == [rule_1, rule_2, rule_3]


@mock.patch('ddtrace.sampler.RateByServiceSampler.sample')
def test_datadog_sampler_sample_no_rules(mock_sample, dummy_tracer):
    sampler = DatadogSampler()
    span = create_span(tracer=dummy_tracer)

    # Default RateByServiceSampler() is applied
    #   No rules configured
    #   No global rate limit
    #   No rate limit configured
    # RateByServiceSampler.sample(span) returns True
    mock_sample.return_value = True
    assert sampler.sample(span) is True
    assert span._context.sampling_priority is AUTO_KEEP
    assert span.sampled is True

    span = create_span(tracer=dummy_tracer)

    # Default RateByServiceSampler() is applied
    #   No rules configured
    #   No global rate limit
    #   No rate limit configured
    # RateByServiceSampler.sample(span) returns False
    mock_sample.return_value = False
    assert sampler.sample(span) is False
    assert span._context.sampling_priority is AUTO_REJECT
    assert span.sampled is False


@mock.patch('ddtrace.internal.rate_limiter.RateLimiter.is_allowed')
def test_datadog_sampler_sample_rules(mock_is_allowed, dummy_tracer):
    # Do not let the limiter get in the way of our test
    mock_is_allowed.return_value = True

    rules = [
        mock.Mock(spec=SamplingRule),
        mock.Mock(spec=SamplingRule),
        mock.Mock(spec=SamplingRule),
    ]
    sampler = DatadogSampler(rules=rules)

    # Reset all of our mocks
    @contextlib.contextmanager
    def reset_mocks():
        def reset():
            mock_is_allowed.reset_mock()
            for rule in rules:
                rule.reset_mock()
                rule.sample_rate = 0.5

            default_rule = SamplingRule(sample_rate=1.0)
            sampler.default_sampler = mock.Mock(spec=SamplingRule, wraps=default_rule)
            # Mock has lots of problems with mocking/wrapping over class properties
            sampler.default_sampler.sample_rate = default_rule.sample_rate

        reset()  # Reset before, just in case
        try:
            yield
        finally:
            reset()  # Must reset after

    # No rules want to sample
    #   It is allowed because of default rate sampler
    #   All rules SamplingRule.matches are called
    #   No calls to SamplingRule.sample happen
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)
        for rule in rules:
            rule.matches.return_value = False

        assert sampler.sample(span) is True
        assert span._context.sampling_priority is AUTO_KEEP
        assert span.sampled is True
        mock_is_allowed.assert_called_once_with()
        for rule in rules:
            rule.matches.assert_called_once_with(span)
            rule.sample.assert_not_called()
        sampler.default_sampler.matches.assert_not_called()
        sampler.default_sampler.sample.assert_called_once_with(span)
        assert_sampling_decision_tags(span, rule=1.0, limit=1.0)

    # One rule thinks it should be sampled
    #   All following rule's SamplingRule.matches are not called
    #   It goes through limiter
    #   It is allowed
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)

        rules[1].matches.return_value = True
        rules[1].sample.return_value = True

        assert sampler.sample(span) is True
        assert span._context.sampling_priority is AUTO_KEEP
        assert span.sampled is True
        mock_is_allowed.assert_called_once_with()
        sampler.default_sampler.sample.assert_not_called()
        assert_sampling_decision_tags(span, rule=0.5, limit=1.0)

        rules[0].matches.assert_called_once_with(span)
        rules[0].sample.assert_not_called()

        rules[1].matches.assert_called_once_with(span)
        rules[1].sample.assert_called_once_with(span)

        rules[2].matches.assert_not_called()
        rules[2].sample.assert_not_called()

    # All rules think it should be sampled
    #   The first rule's SamplingRule.matches is called
    #   It goes through limiter
    #   It is allowed
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)

        for rule in rules:
            rule.matches.return_value = True
        rules[0].sample.return_value = True

        assert sampler.sample(span) is True
        assert span._context.sampling_priority is AUTO_KEEP
        assert span.sampled is True
        mock_is_allowed.assert_called_once_with()
        sampler.default_sampler.sample.assert_not_called()
        assert_sampling_decision_tags(span, rule=0.5, limit=1.0)

        rules[0].matches.assert_called_once_with(span)
        rules[0].sample.assert_called_once_with(span)
        for rule in rules[1:]:
            rule.matches.assert_not_called()
            rule.sample.assert_not_called()

    # Rule matches but does not think it should be sampled
    #   The rule's SamplingRule.matches is called
    #   The rule's SamplingRule.sample is called
    #   Rate limiter is not called
    #   The span is rejected
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)

        rules[0].matches.return_value = False
        rules[2].matches.return_value = False

        rules[1].matches.return_value = True
        rules[1].sample.return_value = False

        assert sampler.sample(span) is False
        assert span._context.sampling_priority is AUTO_REJECT
        assert span.sampled is False
        mock_is_allowed.assert_not_called()
        sampler.default_sampler.sample.assert_not_called()
        assert_sampling_decision_tags(span, rule=0.5)

        rules[0].matches.assert_called_once_with(span)
        rules[0].sample.assert_not_called()

        rules[1].matches.assert_called_once_with(span)
        rules[1].sample.assert_called_once_with(span)

        rules[2].matches.assert_not_called()
        rules[2].sample.assert_not_called()

    # No rules match and RateByServiceSampler is used
    #   All rules SamplingRule.matches are called
    #   Priority sampler's `sample` method is called
    #   Result of priority sampler is returned
    #   Rate limiter is not called
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)

        # Configure mock priority sampler
        priority_sampler = RateByServiceSampler()
        sampler.default_sampler = mock.Mock(spec=RateByServiceSampler, wraps=priority_sampler)

        for rule in rules:
            rule.matches.return_value = False
            rule.sample.return_value = False

        assert sampler.sample(span) is True
        assert span._context.sampling_priority is AUTO_KEEP
        assert span.sampled is True
        mock_is_allowed.assert_not_called()
        sampler.default_sampler.sample.assert_called_once_with(span)
        assert_sampling_decision_tags(span, agent=1)

        [r.matches.assert_called_once_with(span) for r in rules]
        [r.sample.assert_not_called() for r in rules]

    # No rules match and priority sampler is defined
    #   All rules SamplingRule.matches are called
    #   Priority sampler's `sample` method is called
    #   Result of priority sampler is returned
    #   Rate limiter is not called
    with reset_mocks():
        span = create_span(tracer=dummy_tracer)

        # Configure mock priority sampler
        priority_sampler = RateByServiceSampler()
        for rate_sampler in priority_sampler._by_service_samplers.values():
            rate_sampler.set_sample_rate(0)

        sampler.default_sampler = mock.Mock(spec=RateByServiceSampler, wraps=priority_sampler)

        for rule in rules:
            rule.matches.return_value = False
            rule.sample.return_value = False

        assert sampler.sample(span) is False
        assert span._context.sampling_priority is AUTO_REJECT
        assert span.sampled is False
        mock_is_allowed.assert_not_called()
        sampler.default_sampler.sample.assert_called_once_with(span)
        assert_sampling_decision_tags(span, agent=0)

        [r.matches.assert_called_once_with(span) for r in rules]
        [r.sample.assert_not_called() for r in rules]


def test_datadog_sampler_tracer(dummy_tracer):
    rule = SamplingRule(sample_rate=1.0, name='test.span')
    rule_spy = mock.Mock(spec=rule, wraps=rule)
    rule_spy.sample_rate = rule.sample_rate

    sampler = DatadogSampler(rules=[rule_spy])
    limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
    sampler.limiter = limiter_spy
    sampler_spy = mock.Mock(spec=sampler, wraps=sampler)

    dummy_tracer.configure(sampler=sampler_spy)

    assert dummy_tracer.sampler is sampler_spy

    with dummy_tracer.trace('test.span') as span:
        # Assert all of our expected functions were called
        sampler_spy.sample.assert_called_once_with(span)
        rule_spy.matches.assert_called_once_with(span)
        rule_spy.sample.assert_called_once_with(span)
        limiter_spy.is_allowed.assert_called_once_with()

        # It must always mark it as sampled
        assert span.sampled is True
        # We know it was sampled because we have a sample rate of 1.0
        assert span._context.sampling_priority is AUTO_KEEP
        assert_sampling_decision_tags(span, rule=1.0)


def test_datadog_sampler_tracer_rate_limited(dummy_tracer):
    rule = SamplingRule(sample_rate=1.0, name='test.span')
    rule_spy = mock.Mock(spec=rule, wraps=rule)
    rule_spy.sample_rate = rule.sample_rate

    sampler = DatadogSampler(rules=[rule_spy])
    limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
    limiter_spy.is_allowed.return_value = False  # Have the limiter deny the span
    sampler.limiter = limiter_spy
    sampler_spy = mock.Mock(spec=sampler, wraps=sampler)

    dummy_tracer.configure(sampler=sampler_spy)

    assert dummy_tracer.sampler is sampler_spy

    with dummy_tracer.trace('test.span') as span:
        # Assert all of our expected functions were called
        sampler_spy.sample.assert_called_once_with(span)
        rule_spy.matches.assert_called_once_with(span)
        rule_spy.sample.assert_called_once_with(span)
        limiter_spy.is_allowed.assert_called_once_with()

        # We must always mark the span as sampled
        assert span.sampled is True
        assert span._context.sampling_priority is AUTO_REJECT
        assert_sampling_decision_tags(span, rule=1.0, limit=None)


def test_datadog_sampler_tracer_rate_0(dummy_tracer):
    rule = SamplingRule(sample_rate=0, name='test.span')  # Sample rate of 0 means never sample
    rule_spy = mock.Mock(spec=rule, wraps=rule)
    rule_spy.sample_rate = rule.sample_rate

    sampler = DatadogSampler(rules=[rule_spy])
    limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
    sampler.limiter = limiter_spy
    sampler_spy = mock.Mock(spec=sampler, wraps=sampler)

    dummy_tracer.configure(sampler=sampler_spy)

    assert dummy_tracer.sampler is sampler_spy

    with dummy_tracer.trace('test.span') as span:
        # Assert all of our expected functions were called
        sampler_spy.sample.assert_called_once_with(span)
        rule_spy.matches.assert_called_once_with(span)
        rule_spy.sample.assert_called_once_with(span)
        limiter_spy.is_allowed.assert_not_called()

        # It must always mark it as sampled
        assert span.sampled is True
        # We know it was not sampled because we have a sample rate of 0.0
        assert span._context.sampling_priority is AUTO_REJECT
        assert_sampling_decision_tags(span, rule=0)


def test_datadog_sampler_tracer_child(dummy_tracer):
    rule = SamplingRule(sample_rate=1.0)  # No rules means it gets applied to every span
    rule_spy = mock.Mock(spec=rule, wraps=rule)
    rule_spy.sample_rate = rule.sample_rate

    sampler = DatadogSampler(rules=[rule_spy])
    limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
    sampler.limiter = limiter_spy
    sampler_spy = mock.Mock(spec=sampler, wraps=sampler)

    dummy_tracer.configure(sampler=sampler_spy)

    assert dummy_tracer.sampler is sampler_spy

    with dummy_tracer.trace('parent.span') as parent:
        with dummy_tracer.trace('child.span') as child:
            # Assert all of our expected functions were called
            # DEV: `assert_called_once_with` ensures we didn't also call with the child span
            sampler_spy.sample.assert_called_once_with(parent)
            rule_spy.matches.assert_called_once_with(parent)
            rule_spy.sample.assert_called_once_with(parent)
            limiter_spy.is_allowed.assert_called_once_with()

            # We know it was sampled because we have a sample rate of 1.0
            assert parent.sampled is True
            assert parent._context.sampling_priority is AUTO_KEEP
            assert_sampling_decision_tags(parent, rule=1.0)

            assert child.sampled is True
            assert child._parent is parent
            assert child._context.sampling_priority is AUTO_KEEP


def test_datadog_sampler_tracer_start_span(dummy_tracer):
    rule = SamplingRule(sample_rate=1.0)  # No rules means it gets applied to every span
    rule_spy = mock.Mock(spec=rule, wraps=rule)
    rule_spy.sample_rate = rule.sample_rate

    sampler = DatadogSampler(rules=[rule_spy])
    limiter_spy = mock.Mock(spec=sampler.limiter, wraps=sampler.limiter)
    sampler.limiter = limiter_spy
    sampler_spy = mock.Mock(spec=sampler, wraps=sampler)

    dummy_tracer.configure(sampler=sampler_spy)

    assert dummy_tracer.sampler is sampler_spy

    span = dummy_tracer.start_span('test.span')

    # Assert all of our expected functions were called
    sampler_spy.sample.assert_called_once_with(span)
    rule_spy.matches.assert_called_once_with(span)
    rule_spy.sample.assert_called_once_with(span)
    limiter_spy.is_allowed.assert_called_once_with()

    # It must always mark it as sampled
    assert span.sampled is True
    # We know it was sampled because we have a sample rate of 1.0
    assert span._context.sampling_priority is AUTO_KEEP
    assert_sampling_decision_tags(span, rule=1.0)


def test_datadog_sampler_update_rate_by_service_sample_rates(dummy_tracer):
    cases = [
        {
            'service:,env:': 1,
        },
        {
            'service:,env:': 1,
            'service:mcnulty,env:dev': 0.33,
            'service:postgres,env:dev': 0.7,
        },
        {
            'service:,env:': 1,
            'service:mcnulty,env:dev': 0.25,
            'service:postgres,env:dev': 0.5,
            'service:redis,env:prod': 0.75,
        },
    ]

    # By default sampler sets it's default sampler to RateByServiceSampler
    sampler = DatadogSampler()
    for case in cases:
        sampler.update_rate_by_service_sample_rates(case)
        rates = {}
        for k, v in iteritems(sampler.default_sampler._by_service_samplers):
            rates[k] = v.sample_rate
        assert case == rates, '%s != %s' % (case, rates)

    # It's important to also test in reverse mode for we want to make sure key deletion
    # works as well as key insertion (and doing this both ways ensures we trigger both cases)
    cases.reverse()
    for case in cases:
        sampler.update_rate_by_service_sample_rates(case)
        rates = {}
        for k, v in iteritems(sampler.default_sampler._by_service_samplers):
            rates[k] = v.sample_rate
        assert case == rates, '%s != %s' % (case, rates)