def test_opentelemetry_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 = OpenTelemetrySampler(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)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
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()
# We know it was sampled because we have a sample rate of 1.0
assert span.sampled is True
assert span._context.sampling_priority is AUTO_KEEP
assert_sampling_decision_tags(span, rule=1.0)
def test_opentelemetry_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 = OpenTelemetrySampler(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)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
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_opentelemetry_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 = OpenTelemetrySampler(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)
# TODO: Remove `priority_sampling=False` when we remove fallback
dummy_tracer.configure(sampler=sampler_spy, priority_sampling=False)
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()
# We know it was not sampled because we have a sample rate of 0.0
assert span.sampled is False
assert span._context.sampling_priority is AUTO_REJECT
assert_sampling_decision_tags(span, rule=0)
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_init():
# No args
sampler = OpenTelemetrySampler()
assert sampler.rules == []
assert isinstance(sampler.limiter, RateLimiter)
assert sampler.limiter.rate_limit == OpenTelemetrySampler.DEFAULT_RATE_LIMIT
# With rules
rule = SamplingRule(sample_rate=1)
sampler = OpenTelemetrySampler(rules=[rule])
assert sampler.rules == [rule]
assert sampler.limiter.rate_limit == OpenTelemetrySampler.DEFAULT_RATE_LIMIT
# With rate limit
sampler = OpenTelemetrySampler(rate_limit=10)
assert sampler.limiter.rate_limit == 10
# Invalid rules
for val in (None, True, False, object(), 1, Exception()):
with pytest.raises(TypeError):
OpenTelemetrySampler(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 = OpenTelemetrySampler(rules=[rule_1, rule_2, rule_3])
assert sampler.rules == [rule_1, rule_2, rule_3]
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