@@ -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
+ self ,
33
+ inject : mock .MagicMock ,
34
+ set_span_in_context : mock .MagicMock ,
35
+ enrich_span : mock .MagicMock ,
36
+ extract_send_partition : mock .MagicMock ,
37
+ extract_bootstrap_servers : mock .MagicMock ,
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 (
32
75
self ,
33
76
inject : mock .MagicMock ,
34
77
set_span_in_context : mock .MagicMock ,
35
78
enrich_span : mock .MagicMock ,
36
79
extract_send_partition : mock .MagicMock ,
37
80
extract_bootstrap_servers : mock .MagicMock ,
38
- ):
81
+ ) -> None :
39
82
tracer = mock .MagicMock ()
40
83
produce_hook = mock .MagicMock ()
41
84
original_send_callback = mock .MagicMock ()
0 commit comments