21
21
from logging import ERROR , WARNING
22
22
from random import randint
23
23
from time import time_ns
24
- from typing import Optional
24
+ from typing import Optional , Sequence
25
25
from unittest import mock
26
26
27
27
from opentelemetry import trace as trace_api
39
39
OTEL_TRACES_SAMPLER ,
40
40
OTEL_TRACES_SAMPLER_ARG ,
41
41
)
42
- from opentelemetry .sdk .trace import Resource , sampling
42
+ from opentelemetry .sdk .trace import Resource
43
43
from opentelemetry .sdk .trace .id_generator import RandomIdGenerator
44
+ from opentelemetry .sdk .trace .sampling import (
45
+ ALWAYS_OFF ,
46
+ ALWAYS_ON ,
47
+ Decision ,
48
+ ParentBased ,
49
+ Sampler ,
50
+ SamplingResult ,
51
+ StaticSampler ,
52
+ TraceIdRatioBased ,
53
+ )
44
54
from opentelemetry .sdk .util import ns_to_iso_str
45
55
from opentelemetry .sdk .util .instrumentation import InstrumentationInfo
46
56
from opentelemetry .test .spantestutil import (
47
57
get_span_with_dropped_attributes_events_links ,
48
58
new_tracer ,
49
59
)
50
- from opentelemetry .trace import Status , StatusCode
60
+ from opentelemetry .trace import Link , SpanKind , Status , StatusCode
61
+ from opentelemetry .trace .span import TraceState
62
+ from opentelemetry .util .types import Attributes
51
63
52
64
53
65
class TestTracer (unittest .TestCase ):
@@ -139,60 +151,52 @@ def test_tracer_provider_accepts_concurrent_multi_span_processor(self):
139
151
)
140
152
141
153
142
- class CustomSampler (sampling . Sampler ):
154
+ class CustomSampler (Sampler ):
143
155
def __init__ (self ) -> None :
144
- super (). __init__ ()
156
+ pass
145
157
146
158
def get_description (self ) -> str :
147
- return super (). get_description ()
159
+ return "CustomSampler"
148
160
149
161
def should_sample (
150
162
self ,
151
- parent_context ,
152
- trace_id ,
153
- name ,
154
- kind ,
155
- attributes ,
156
- links ,
157
- trace_state ,
158
- ):
159
- return super ().should_sample (
160
- parent_context ,
161
- trace_id ,
162
- name ,
163
- kind ,
164
- attributes ,
165
- links ,
166
- trace_state ,
167
- )
168
-
169
-
170
- class CustomRatioSampler (sampling .TraceIdRatioBased ):
163
+ parent_context : Optional ["Context" ],
164
+ trace_id : int ,
165
+ name : str ,
166
+ kind : SpanKind = None ,
167
+ attributes : Attributes = None ,
168
+ links : Sequence [Link ] = None ,
169
+ trace_state : TraceState = None ,
170
+ ) -> "SamplingResult" :
171
+ return SamplingResult (
172
+ Decision .RECORD_AND_SAMPLE ,
173
+ None ,
174
+ None ,
175
+ )
176
+
177
+
178
+ class CustomRatioSampler (TraceIdRatioBased ):
171
179
def __init__ (self , ratio ):
172
180
self .ratio = ratio
173
181
super ().__init__ (ratio )
174
182
175
183
def get_description (self ) -> str :
176
- return super (). get_description ()
184
+ return "CustomSampler"
177
185
178
186
def should_sample (
179
187
self ,
180
- parent_context ,
181
- trace_id ,
182
- name ,
183
- kind ,
184
- attributes ,
185
- links ,
186
- trace_state ,
187
- ):
188
- return super ().should_sample (
189
- parent_context ,
190
- trace_id ,
191
- name ,
192
- kind ,
193
- attributes ,
194
- links ,
195
- trace_state ,
188
+ parent_context : Optional ["Context" ],
189
+ trace_id : int ,
190
+ name : str ,
191
+ kind : SpanKind = None ,
192
+ attributes : Attributes = None ,
193
+ links : Sequence [Link ] = None ,
194
+ trace_state : TraceState = None ,
195
+ ) -> "SamplingResult" :
196
+ return SamplingResult (
197
+ Decision .RECORD_AND_SAMPLE ,
198
+ None ,
199
+ None ,
196
200
)
197
201
198
202
@@ -248,7 +252,7 @@ def test_default_sampler_type(self):
248
252
self .verify_default_sampler (tracer_provider )
249
253
250
254
def test_sampler_no_sampling (self ):
251
- tracer_provider = trace .TracerProvider (sampling . ALWAYS_OFF )
255
+ tracer_provider = trace .TracerProvider (ALWAYS_OFF )
252
256
tracer = tracer_provider .get_tracer (__name__ )
253
257
254
258
# Check that the default tracer creates no-op spans if the sampler
@@ -272,10 +276,8 @@ def test_sampler_with_env(self):
272
276
# pylint: disable=protected-access
273
277
reload (trace )
274
278
tracer_provider = trace .TracerProvider ()
275
- self .assertIsInstance (tracer_provider .sampler , sampling .StaticSampler )
276
- self .assertEqual (
277
- tracer_provider .sampler ._decision , sampling .Decision .DROP
278
- )
279
+ self .assertIsInstance (tracer_provider .sampler , StaticSampler )
280
+ self .assertEqual (tracer_provider .sampler ._decision , Decision .DROP )
279
281
280
282
tracer = tracer_provider .get_tracer (__name__ )
281
283
@@ -294,7 +296,7 @@ def test_ratio_sampler_with_env(self):
294
296
# pylint: disable=protected-access
295
297
reload (trace )
296
298
tracer_provider = trace .TracerProvider ()
297
- self .assertIsInstance (tracer_provider .sampler , sampling . ParentBased )
299
+ self .assertIsInstance (tracer_provider .sampler , ParentBased )
298
300
self .assertEqual (tracer_provider .sampler ._root .rate , 0.25 )
299
301
300
302
@mock .patch .dict (
@@ -453,9 +455,9 @@ def test_custom_ratio_sampler_with_env_multiple_entry_points(
453
455
self .assertIsInstance (tracer_provider .sampler , CustomSampler )
454
456
455
457
def verify_default_sampler (self , tracer_provider ):
456
- self .assertIsInstance (tracer_provider .sampler , sampling . ParentBased )
458
+ self .assertIsInstance (tracer_provider .sampler , ParentBased )
457
459
# pylint: disable=protected-access
458
- self .assertEqual (tracer_provider .sampler ._root , sampling . ALWAYS_ON )
460
+ self .assertEqual (tracer_provider .sampler ._root , ALWAYS_ON )
459
461
460
462
461
463
class TestSpanCreation (unittest .TestCase ):
@@ -950,7 +952,7 @@ def test_sampling_attributes(self):
950
952
"attr-in-both" : "decision-attr" ,
951
953
}
952
954
tracer_provider = trace .TracerProvider (
953
- sampling . StaticSampler (sampling . Decision .RECORD_AND_SAMPLE )
955
+ StaticSampler (Decision .RECORD_AND_SAMPLE )
954
956
)
955
957
956
958
self .tracer = tracer_provider .get_tracer (__name__ )
0 commit comments