22
22
from opentelemetry .instrumentation .sqlalchemy import SQLAlchemyInstrumentor
23
23
from opentelemetry .sdk .resources import Resource
24
24
from opentelemetry .sdk .trace import TracerProvider , export
25
+ from opentelemetry .semconv .trace import SpanAttributes
25
26
from opentelemetry .test .test_base import TestBase
26
27
27
28
@@ -128,11 +129,12 @@ async def run():
128
129
def test_not_recording (self ):
129
130
mock_tracer = mock .Mock ()
130
131
mock_span = mock .Mock ()
132
+ mock_context = mock .Mock ()
131
133
mock_span .is_recording .return_value = False
132
- mock_span .__enter__ = mock .Mock (return_value = ( mock . Mock (), None ) )
133
- mock_span .__exit__ = mock .Mock (return_value = None )
134
- mock_tracer .start_span .return_value = mock_span
135
- mock_tracer .start_as_current_span .return_value = mock_span
134
+ mock_context .__enter__ = mock .Mock (return_value = mock_span )
135
+ mock_context .__exit__ = mock .Mock (return_value = None )
136
+ mock_tracer .start_span .return_value = mock_context
137
+ mock_tracer .start_as_current_span .return_value = mock_context
136
138
with mock .patch ("opentelemetry.trace.get_tracer" ) as tracer :
137
139
tracer .return_value = mock_tracer
138
140
engine = create_engine ("sqlite:///:memory:" )
@@ -159,6 +161,12 @@ def test_create_engine_wrapper(self):
159
161
self .assertEqual (len (spans ), 2 )
160
162
# first span - the connection to the db
161
163
self .assertEqual (spans [0 ].name , "connect" )
164
+ self .assertEqual (
165
+ spans [0 ].attributes [SpanAttributes .DB_NAME ], ":memory:"
166
+ )
167
+ self .assertEqual (
168
+ spans [0 ].attributes [SpanAttributes .DB_SYSTEM ], "sqlite"
169
+ )
162
170
self .assertEqual (spans [0 ].kind , trace .SpanKind .CLIENT )
163
171
# second span - the query
164
172
self .assertEqual (spans [1 ].name , "SELECT :memory:" )
@@ -217,6 +225,12 @@ async def run():
217
225
self .assertEqual (len (spans ), 2 )
218
226
# first span - the connection to the db
219
227
self .assertEqual (spans [0 ].name , "connect" )
228
+ self .assertEqual (
229
+ spans [0 ].attributes [SpanAttributes .DB_NAME ], ":memory:"
230
+ )
231
+ self .assertEqual (
232
+ spans [0 ].attributes [SpanAttributes .DB_SYSTEM ], "sqlite"
233
+ )
220
234
self .assertEqual (spans [0 ].kind , trace .SpanKind .CLIENT )
221
235
# second span - the query
222
236
self .assertEqual (spans [1 ].name , "SELECT :memory:" )
0 commit comments