@@ -28,14 +28,57 @@ def setUp(self) -> None:
28
28
@mock .patch ("opentelemetry.instrumentation.kafka.utils._enrich_span" )
29
29
@mock .patch ("opentelemetry.trace.set_span_in_context" )
30
30
@mock .patch ("opentelemetry.propagate.inject" )
31
- def test_wrap_send (
31
+ def test_wrap_send_with_topic_as_arg (
32
32
self ,
33
33
inject : mock .MagicMock ,
34
34
set_span_in_context : mock .MagicMock ,
35
35
enrich_span : mock .MagicMock ,
36
36
extract_send_partition : mock .MagicMock ,
37
37
extract_bootstrap_servers : mock .MagicMock ,
38
- ):
38
+ ) -> None :
39
+ self .wrap_send_helper (
40
+ inject ,
41
+ set_span_in_context ,
42
+ enrich_span ,
43
+ extract_send_partition ,
44
+ extract_bootstrap_servers ,
45
+ )
46
+
47
+ @mock .patch (
48
+ "opentelemetry.instrumentation.kafka.utils.KafkaPropertiesExtractor.extract_bootstrap_servers"
49
+ )
50
+ @mock .patch (
51
+ "opentelemetry.instrumentation.kafka.utils.KafkaPropertiesExtractor.extract_send_partition"
52
+ )
53
+ @mock .patch ("opentelemetry.instrumentation.kafka.utils._enrich_span" )
54
+ @mock .patch ("opentelemetry.trace.set_span_in_context" )
55
+ @mock .patch ("opentelemetry.propagate.inject" )
56
+ def test_wrap_send_with_topic_as_kwarg (
57
+ self ,
58
+ inject : mock .MagicMock ,
59
+ set_span_in_context : mock .MagicMock ,
60
+ enrich_span : mock .MagicMock ,
61
+ extract_send_partition : mock .MagicMock ,
62
+ extract_bootstrap_servers : mock .MagicMock ,
63
+ ) -> None :
64
+ self .args = []
65
+ self .kwargs ["topic" ] = self .topic_name
66
+ self .wrap_send_helper (
67
+ inject ,
68
+ set_span_in_context ,
69
+ enrich_span ,
70
+ extract_send_partition ,
71
+ extract_bootstrap_servers ,
72
+ )
73
+
74
+ def wrap_send_helper (
75
+ self ,
76
+ inject : mock .MagicMock ,
77
+ set_span_in_context : mock .MagicMock ,
78
+ enrich_span : mock .MagicMock ,
79
+ extract_send_partition : mock .MagicMock ,
80
+ extract_bootstrap_servers : mock .MagicMock ,
81
+ ) -> None :
39
82
tracer = mock .MagicMock ()
40
83
produce_hook = mock .MagicMock ()
41
84
original_send_callback = mock .MagicMock ()
0 commit comments