2
2
import uuid
3
3
4
4
import pytest
5
- from opentelemetry .trace import SpanKind , get_tracer
6
- from opentelemetry .trace .span import Span
7
5
8
6
from opentelemetry .instrumentation .langchain .span_manager import (
9
7
SpanManager ,
10
8
_SpanState ,
11
9
)
10
+ from opentelemetry .trace import SpanKind , get_tracer
11
+ from opentelemetry .trace .span import Span
12
12
13
13
14
14
class TestSpanManager :
@@ -28,8 +28,9 @@ def handler(self, tracer):
28
28
(uuid .uuid4 (), True ), # Parent in spans
29
29
],
30
30
)
31
-
32
- def test_create_span (self , handler , tracer , parent_run_id , parent_in_spans ):
31
+ def test_create_span (
32
+ self , handler , tracer , parent_run_id , parent_in_spans
33
+ ):
33
34
# Arrange
34
35
run_id = uuid .uuid4 ()
35
36
span_name = "test_span"
@@ -40,17 +41,25 @@ def test_create_span(self, handler, tracer, parent_run_id, parent_in_spans):
40
41
# Setup parent if needed
41
42
if parent_run_id is not None and parent_in_spans :
42
43
parent_mock_span = unittest .mock .Mock (spec = Span )
43
- handler .spans [parent_run_id ] = _SpanState (span = parent_mock_span , context = None )
44
+ handler .spans [parent_run_id ] = _SpanState (
45
+ span = parent_mock_span , context = None
46
+ )
44
47
45
- with unittest .mock .patch .object (
46
- tracer , "start_span" , return_value = mock_span
47
- ) as mock_start_span , unittest .mock .patch (
48
- "opentelemetry.instrumentation.langchain.span_manager.set_span_in_context"
49
- ) as mock_set_span_in_context , unittest .mock .patch (
50
- "opentelemetry.instrumentation.langchain.span_manager.get_current"
51
- ) as mock_get_current :
48
+ with (
49
+ unittest .mock .patch .object (
50
+ tracer , "start_span" , return_value = mock_span
51
+ ) as mock_start_span ,
52
+ unittest .mock .patch (
53
+ "opentelemetry.instrumentation.langchain.span_manager.set_span_in_context"
54
+ ) as mock_set_span_in_context ,
55
+ unittest .mock .patch (
56
+ "opentelemetry.instrumentation.langchain.span_manager.get_current"
57
+ ) as mock_get_current ,
58
+ ):
52
59
# Act
53
- result = handler .create_span (run_id , parent_run_id , span_name , kind )
60
+ result = handler .create_span (
61
+ run_id , parent_run_id , span_name , kind
62
+ )
54
63
55
64
# Assert
56
65
assert result == mock_span
@@ -59,11 +68,17 @@ def test_create_span(self, handler, tracer, parent_run_id, parent_in_spans):
59
68
60
69
# Verify parent-child relationship
61
70
if parent_run_id is not None and parent_in_spans :
62
- mock_set_span_in_context .assert_called_once_with (handler .spans [parent_run_id ].span )
71
+ mock_set_span_in_context .assert_called_once_with (
72
+ handler .spans [parent_run_id ].span
73
+ )
63
74
mock_start_span .assert_called_once_with (
64
- name = span_name , kind = kind , context = mock_set_span_in_context .return_value
75
+ name = span_name ,
76
+ kind = kind ,
77
+ context = mock_set_span_in_context .return_value ,
65
78
)
66
79
assert run_id in handler .spans [parent_run_id ].children
67
80
else :
68
- mock_start_span .assert_called_once_with (name = span_name , kind = kind )
69
- mock_set_span_in_context .assert_not_called ()
81
+ mock_start_span .assert_called_once_with (
82
+ name = span_name , kind = kind
83
+ )
84
+ mock_set_span_in_context .assert_not_called ()
0 commit comments