@@ -26,6 +26,8 @@ def get_as_list(dict_object, key):
26
26
def mock_inject (name , value = "data" ):
27
27
def wrapped (setter , carrier = None , context = None ):
28
28
carrier [name ] = value
29
+ setter (None , "inject_field_{}_0" .format (name ), None )
30
+ setter (None , "inject_field_{}_1" .format (name ), None )
29
31
30
32
return wrapped
31
33
@@ -39,18 +41,29 @@ def wrapped(getter, carrier=None, context=None):
39
41
return wrapped
40
42
41
43
44
+ def mock_fields (name ):
45
+ return {
46
+ "inject_field_{}_0" .format (name ), "inject_field_{}_1" .format (name )
47
+ }
48
+
49
+
42
50
class TestCompositePropagator (unittest .TestCase ):
43
51
@classmethod
44
52
def setUpClass (cls ):
45
53
cls .mock_propagator_0 = Mock (
46
- inject = mock_inject ("mock-0" ), extract = mock_extract ("mock-0" )
54
+ inject = mock_inject ("mock-0" ),
55
+ extract = mock_extract ("mock-0" ),
56
+ fields = mock_fields ("mock-0" ),
47
57
)
48
58
cls .mock_propagator_1 = Mock (
49
- inject = mock_inject ("mock-1" ), extract = mock_extract ("mock-1" )
59
+ inject = mock_inject ("mock-1" ),
60
+ extract = mock_extract ("mock-1" ),
61
+ fields = mock_fields ("mock-1" ),
50
62
)
51
63
cls .mock_propagator_2 = Mock (
52
64
inject = mock_inject ("mock-0" , value = "data2" ),
53
65
extract = mock_extract ("mock-0" , value = "context2" ),
66
+ fields = mock_fields ("mock-0" ),
54
67
)
55
68
56
69
def test_no_propagators (self ):
@@ -105,3 +118,23 @@ def test_multiple_propagators_same_key(self):
105
118
get_as_list , carrier = new_carrier , context = {}
106
119
)
107
120
self .assertEqual (context , {"mock-0" : "context2" })
121
+
122
+ def test_fields (self ):
123
+ propagator = CompositeHTTPPropagator (
124
+ [
125
+ self .mock_propagator_0 ,
126
+ self .mock_propagator_1 ,
127
+ self .mock_propagator_2
128
+ ]
129
+ )
130
+
131
+ mock_set_in_carrier = Mock ()
132
+
133
+ propagator .inject (mock_set_in_carrier , {})
134
+
135
+ inject_fields = set ()
136
+
137
+ for mock_call in mock_set_in_carrier .mock_calls :
138
+ inject_fields .add (mock_call .args [1 ])
139
+
140
+ self .assertEqual (inject_fields , propagator .fields )
0 commit comments