2121from logging import ERROR , WARNING
2222from random import randint
2323from time import time_ns
24- from typing import Optional
24+ from typing import Optional , Sequence
2525from unittest import mock
2626
2727from opentelemetry import trace as trace_api
3939 OTEL_TRACES_SAMPLER ,
4040 OTEL_TRACES_SAMPLER_ARG ,
4141)
42- from opentelemetry .sdk .trace import Resource , sampling
42+ from opentelemetry .sdk .trace import Resource
4343from 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+ )
4454from opentelemetry .sdk .util import ns_to_iso_str
4555from opentelemetry .sdk .util .instrumentation import InstrumentationInfo
4656from opentelemetry .test .spantestutil import (
4757 get_span_with_dropped_attributes_events_links ,
4858 new_tracer ,
4959)
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
5163
5264
5365class TestTracer (unittest .TestCase ):
@@ -139,60 +151,52 @@ def test_tracer_provider_accepts_concurrent_multi_span_processor(self):
139151 )
140152
141153
142- class CustomSampler (sampling . Sampler ):
154+ class CustomSampler (Sampler ):
143155 def __init__ (self ) -> None :
144- super (). __init__ ()
156+ pass
145157
146158 def get_description (self ) -> str :
147- return super (). get_description ()
159+ return "CustomSampler"
148160
149161 def should_sample (
150162 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 ):
171179 def __init__ (self , ratio ):
172180 self .ratio = ratio
173181 super ().__init__ (ratio )
174182
175183 def get_description (self ) -> str :
176- return super (). get_description ()
184+ return "CustomSampler"
177185
178186 def should_sample (
179187 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 ,
196200 )
197201
198202
@@ -248,7 +252,7 @@ def test_default_sampler_type(self):
248252 self .verify_default_sampler (tracer_provider )
249253
250254 def test_sampler_no_sampling (self ):
251- tracer_provider = trace .TracerProvider (sampling . ALWAYS_OFF )
255+ tracer_provider = trace .TracerProvider (ALWAYS_OFF )
252256 tracer = tracer_provider .get_tracer (__name__ )
253257
254258 # Check that the default tracer creates no-op spans if the sampler
@@ -272,10 +276,8 @@ def test_sampler_with_env(self):
272276 # pylint: disable=protected-access
273277 reload (trace )
274278 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 )
279281
280282 tracer = tracer_provider .get_tracer (__name__ )
281283
@@ -294,7 +296,7 @@ def test_ratio_sampler_with_env(self):
294296 # pylint: disable=protected-access
295297 reload (trace )
296298 tracer_provider = trace .TracerProvider ()
297- self .assertIsInstance (tracer_provider .sampler , sampling . ParentBased )
299+ self .assertIsInstance (tracer_provider .sampler , ParentBased )
298300 self .assertEqual (tracer_provider .sampler ._root .rate , 0.25 )
299301
300302 @mock .patch .dict (
@@ -453,9 +455,9 @@ def test_custom_ratio_sampler_with_env_multiple_entry_points(
453455 self .assertIsInstance (tracer_provider .sampler , CustomSampler )
454456
455457 def verify_default_sampler (self , tracer_provider ):
456- self .assertIsInstance (tracer_provider .sampler , sampling . ParentBased )
458+ self .assertIsInstance (tracer_provider .sampler , ParentBased )
457459 # pylint: disable=protected-access
458- self .assertEqual (tracer_provider .sampler ._root , sampling . ALWAYS_ON )
460+ self .assertEqual (tracer_provider .sampler ._root , ALWAYS_ON )
459461
460462
461463class TestSpanCreation (unittest .TestCase ):
@@ -950,7 +952,7 @@ def test_sampling_attributes(self):
950952 "attr-in-both" : "decision-attr" ,
951953 }
952954 tracer_provider = trace .TracerProvider (
953- sampling . StaticSampler (sampling . Decision .RECORD_AND_SAMPLE )
955+ StaticSampler (Decision .RECORD_AND_SAMPLE )
954956 )
955957
956958 self .tracer = tracer_provider .get_tracer (__name__ )
0 commit comments