@@ -26,6 +26,8 @@ def get_as_list(dict_object, key):
2626def mock_inject (name , value = "data" ):
2727 def wrapped (setter , carrier = None , context = None ):
2828 carrier [name ] = value
29+ setter (None , "inject_field_{}_0" .format (name ), None )
30+ setter (None , "inject_field_{}_1" .format (name ), None )
2931
3032 return wrapped
3133
@@ -39,18 +41,29 @@ def wrapped(getter, carrier=None, context=None):
3941 return wrapped
4042
4143
44+ def mock_fields (name ):
45+ return {
46+ "inject_field_{}_0" .format (name ), "inject_field_{}_1" .format (name )
47+ }
48+
49+
4250class TestCompositePropagator (unittest .TestCase ):
4351 @classmethod
4452 def setUpClass (cls ):
4553 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" ),
4757 )
4858 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" ),
5062 )
5163 cls .mock_propagator_2 = Mock (
5264 inject = mock_inject ("mock-0" , value = "data2" ),
5365 extract = mock_extract ("mock-0" , value = "context2" ),
66+ fields = mock_fields ("mock-0" ),
5467 )
5568
5669 def test_no_propagators (self ):
@@ -105,3 +118,23 @@ def test_multiple_propagators_same_key(self):
105118 get_as_list , carrier = new_carrier , context = {}
106119 )
107120 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