Use String strategyType instead of int (#172)

Author: black-adder

jaeger_client/sampler.py

@@ -33,9 +33,6 @@
 from .metrics import Metrics, LegacyMetricsFactory
 from .utils import ErrorReporter
 from .rate_limiter import RateLimiter
-from jaeger_client.thrift_gen.sampling import (
-    SamplingManager
-)
 
 default_logger = logging.getLogger('jaeger_tracing')
 
@@ -55,6 +52,8 @@
 MAX_TRACES_PER_SECOND_STR = 'maxTracesPerSecond'
 RATE_LIMITING_SAMPLING_STR = 'rateLimitingSampling'
 STRATEGY_TYPE_STR = 'strategyType'
+PROBABILISTIC_SAMPLING_STRATEGY = 'PROBABILISTIC'
+RATE_LIMITING_SAMPLING_STRATEGY = 'RATE_LIMITING'
 
 
 class Sampler(object):
@@ -449,10 +448,10 @@ def _update_adaptive_sampler(self, per_operation_strategies):
 
     def _update_rate_limiting_or_probabilistic_sampler(self, response):
         s_type = response.get(STRATEGY_TYPE_STR)
-        if s_type == SamplingManager.SamplingStrategyType.PROBABILISTIC:
+        if s_type == PROBABILISTIC_SAMPLING_STRATEGY:
             sampling_rate = get_sampling_probability(response)
             new_sampler = ProbabilisticSampler(rate=sampling_rate)
-        elif s_type == SamplingManager.SamplingStrategyType.RATE_LIMITING:
+        elif s_type == RATE_LIMITING_SAMPLING_STRATEGY:
             mtps = get_rate_limit(response)
             if 0 <= mtps < 500:
                 new_sampler = RateLimitingSampler(max_traces_per_second=mtps)

tests/test_sampler.py

@@ -355,7 +355,7 @@ def test_sampling_request_callback():
 
     probabilistic_strategy = """
     {
-        "strategyType":0,
+        "strategyType":"PROBABILISTIC",
         "probabilisticSampling":
         {
             "samplingRate":0.002
@@ -374,7 +374,7 @@ def test_sampling_request_callback():
 
     adaptive_sampling_strategy = """
     {
-        "strategyType":0,
+        "strategyType":"PROBABILISTIC",
         "operationSampling":
         {
             "defaultSamplingProbability":0.001,
@@ -427,55 +427,55 @@ def test_sampling_request_callback():
 
 @pytest.mark.parametrize("response,init_sampler,expected_sampler,err_count,err_msg,reference_equivalence", [
     (
-        {"strategyType":0,"probabilisticSampling":{"samplingRate":0.003}},
+        {"strategyType":"PROBABILISTIC","probabilisticSampling":{"samplingRate":0.003}},
         probabilistic_sampler,
         other_probabilistic_sampler,
         0,
         'sampler should update to new probabilistic sampler',
         False,
     ),
     (
-        {"strategyType":0,"probabilisticSampling":{"samplingRate":400}},
+        {"strategyType":"PROBABILISTIC","probabilisticSampling":{"samplingRate":400}},
         probabilistic_sampler,
         probabilistic_sampler,
         1,
         'sampler should remain the same if strategy is invalid',
         True,
     ),
     (
-        {"strategyType":0,"probabilisticSampling":{"samplingRate":0.002}},
+        {"strategyType":"PROBABILISTIC","probabilisticSampling":{"samplingRate":0.002}},
         probabilistic_sampler,
         probabilistic_sampler,
         0,
         'sampler should remain the same with the same strategy',
         True,
     ),
     (
-        {"strategyType":1,"rateLimitingSampling":{"maxTracesPerSecond":10}},
+        {"strategyType":"RATE_LIMITING","rateLimitingSampling":{"maxTracesPerSecond":10}},
         probabilistic_sampler,
         rate_limiting_sampler,
         0,
         'sampler should update to new rate limiting sampler',
         False,
     ),
     (
-        {"strategyType":1,"rateLimitingSampling":{"maxTracesPerSecond":10}},
+        {"strategyType":"RATE_LIMITING","rateLimitingSampling":{"maxTracesPerSecond":10}},
         rate_limiting_sampler,
         rate_limiting_sampler,
         0,
         'sampler should remain the same with the same strategy',
         True,
     ),
     (
-        {"strategyType":1,"rateLimitingSampling":{"maxTracesPerSecond":-10}},
+        {"strategyType":"RATE_LIMITING","rateLimitingSampling":{"maxTracesPerSecond":-10}},
         rate_limiting_sampler,
         rate_limiting_sampler,
         1,
         'sampler should remain the same if strategy is invalid',
         True,
     ),
     (
-        {"strategyType":1,"rateLimitingSampling":{"maxTracesPerSecond":20}},
+        {"strategyType":"RATE_LIMITING","rateLimitingSampling":{"maxTracesPerSecond":20}},
         rate_limiting_sampler,
         other_rate_limiting_sampler,
         0,
@@ -491,7 +491,7 @@ def test_sampling_request_callback():
         True,
     ),
     (
-        {"strategyType":2},
+        {"strategyType":"INVALID_TYPE"},
         rate_limiting_sampler,
         rate_limiting_sampler,
         1,
@@ -533,7 +533,7 @@ def test_update_sampler_adaptive_sampler():
     )
 
     response = {
-        "strategyType":1,
+        "strategyType":"RATE_LIMITING",
         "operationSampling":
         {
             "defaultSamplingProbability":0.001,
@@ -554,7 +554,7 @@ def test_update_sampler_adaptive_sampler():
     assert '%s' % remote_sampler.sampler == 'AdaptiveSampler(0.001000, 2.000000, 10)'
 
     new_response = {
-        "strategyType":1,
+        "strategyType":"RATE_LIMITING",
         "operationSampling":
         {
             "defaultSamplingProbability":0.51,
@@ -574,11 +574,11 @@ def test_update_sampler_adaptive_sampler():
     remote_sampler._update_sampler(new_response)
     assert '%s' % remote_sampler.sampler == 'AdaptiveSampler(0.510000, 3.000000, 10)'
 
-    remote_sampler._update_sampler({"strategyType":0,"probabilisticSampling":{"samplingRate":0.004}})
+    remote_sampler._update_sampler({"strategyType":"PROBABILISTIC","probabilisticSampling":{"samplingRate":0.004}})
     assert '%s' % remote_sampler.sampler == 'ProbabilisticSampler(0.004)', \
         'should not fail going from adaptive sampler to probabilistic sampler'
 
-    remote_sampler._update_sampler({"strategyType":1,"operationSampling":{"defaultSamplingProbability":0.4}})
+    remote_sampler._update_sampler({"strategyType":"RATE_LIMITING","operationSampling":{"defaultSamplingProbability":0.4}})
     assert '%s' % remote_sampler.sampler == 'AdaptiveSampler(0.400000, 0.001667, 10)'
 
     remote_sampler.close()