Example #1
0
 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')
Example #2
0
 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")
Example #3
0
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)
    sampler.default_sampler = mock.Mock(spec=SamplingRule)
    sampler.default_sampler.return_value = True

    # 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
            sampler.default_sampler.reset_mock()
            sampler.default_sampler.sample_rate = 1.0

        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 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
    # TODO: Remove this case when we remove fallback to priority sampling
    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(1)

        spy_sampler = mock.Mock(spec=RateByServiceSampler,
                                wraps=priority_sampler)
        sampler._priority_sampler = spy_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_not_called()
        spy_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]

        # Reset priority sampler property
        sampler._priority_sampler = None

    # 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
    # TODO: Remove this case when we remove fallback to priority sampling
    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)

        spy_sampler = mock.Mock(spec=RateByServiceSampler,
                                wraps=priority_sampler)
        sampler._priority_sampler = spy_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_not_called()
        spy_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]

        # Reset priority sampler property
        sampler._priority_sampler = None