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