def test_opentelemetry_sampler_sample_no_rules(mock_is_allowed, dummy_tracer):
sampler = OpenTelemetrySampler()
span = create_span(tracer=dummy_tracer)
# Default SamplingRule(sample_rate=1.0) is applied
# No priority sampler configured
# No rules configured
# RateLimiter is allowed, it is sampled
mock_is_allowed.return_value = True
assert sampler.sample(span) is True
assert span._context.sampling_priority is AUTO_KEEP
assert span.sampled is True
assert_sampling_decision_tags(span, rule=1.0, limit=1.0)
mock_is_allowed.assert_called_once_with()
mock_is_allowed.reset_mock()
span = create_span(tracer=dummy_tracer)
# Default SamplingRule(sample_rate=1.0) is applied
# No priority sampler configured
# No rules configured
# RateLimit not allowed, it is not sampled
mock_is_allowed.return_value = False
assert sampler.sample(span) is False
assert span._context.sampling_priority is AUTO_REJECT
assert span.sampled is False
# DEV: Is `None` since we only add tag to non-rate limited traces
assert_sampling_decision_tags(span, rule=1.0, limit=None)
mock_is_allowed.assert_called_once_with()
def test_opentelemetry_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 = OpenTelemetrySampler(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