Skip to content

Commit d6b625e

Browse files
tammy-baylis-switammy-baylis-swi
committed
Add tests
1 parent a05acca commit d6b625e

File tree

1 file changed

+181
-0
lines changed

1 file changed

+181
-0
lines changed

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,44 @@ def test_create_engine_wrapper_enable_commenter(self):
209209
self.caplog.records[-2].getMessage(),
210210
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
211211
)
212+
spans = self.memory_exporter.get_finished_spans()
213+
self.assertEqual(len(spans), 2)
214+
# first span is connection to db
215+
self.assertEqual(spans[0].name, "connect")
216+
# second span is query itself
217+
query_span = spans[1]
218+
self.assertEqual(
219+
query_span.attributes[SpanAttributes.DB_STATEMENT],
220+
"SELECT 1;",
221+
)
222+
223+
def test_create_engine_wrapper_enable_commenter_stmt_enabled(self):
224+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
225+
SQLAlchemyInstrumentor().instrument(
226+
enable_commenter=True,
227+
commenter_options={"db_framework": False},
228+
enable_attribute_commenter=True,
229+
)
230+
from sqlalchemy import create_engine # pylint: disable-all
231+
232+
engine = create_engine("sqlite:///:memory:")
233+
cnx = engine.connect()
234+
cnx.execute(text("SELECT 1;")).fetchall()
235+
# sqlcommenter
236+
self.assertRegex(
237+
self.caplog.records[-2].getMessage(),
238+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
239+
)
240+
spans = self.memory_exporter.get_finished_spans()
241+
self.assertEqual(len(spans), 2)
242+
# first span is connection to db
243+
self.assertEqual(spans[0].name, "connect")
244+
# second span is query itself
245+
query_span = spans[1]
246+
self.assertRegex(
247+
query_span.attributes[SpanAttributes.DB_STATEMENT],
248+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
249+
)
212250

213251
def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
214252
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
@@ -229,6 +267,49 @@ def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
229267
self.caplog.records[-2].getMessage(),
230268
r"SELECT 1 /\*db_driver='(.*)'\*/;",
231269
)
270+
spans = self.memory_exporter.get_finished_spans()
271+
self.assertEqual(len(spans), 2)
272+
# first span is connection to db
273+
self.assertEqual(spans[0].name, "connect")
274+
# second span is query itself
275+
query_span = spans[1]
276+
self.assertEqual(
277+
query_span.attributes[SpanAttributes.DB_STATEMENT],
278+
"SELECT 1;",
279+
)
280+
281+
def test_create_engine_wrapper_enable_commenter_stmt_enabled_otel_values_false(
282+
self,
283+
):
284+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
285+
SQLAlchemyInstrumentor().instrument(
286+
enable_commenter=True,
287+
commenter_options={
288+
"db_framework": False,
289+
"opentelemetry_values": False,
290+
},
291+
enable_attribute_commenter=True,
292+
)
293+
from sqlalchemy import create_engine # pylint: disable-all
294+
295+
engine = create_engine("sqlite:///:memory:")
296+
cnx = engine.connect()
297+
cnx.execute(text("SELECT 1;")).fetchall()
298+
# sqlcommenter
299+
self.assertRegex(
300+
self.caplog.records[-2].getMessage(),
301+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
302+
)
303+
spans = self.memory_exporter.get_finished_spans()
304+
self.assertEqual(len(spans), 2)
305+
# first span is connection to db
306+
self.assertEqual(spans[0].name, "connect")
307+
# second span is query itself
308+
query_span = spans[1]
309+
self.assertRegex(
310+
query_span.attributes[SpanAttributes.DB_STATEMENT],
311+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
312+
)
232313

233314
def test_custom_tracer_provider(self):
234315
provider = TracerProvider(
@@ -321,6 +402,55 @@ async def run():
321402
self.caplog.records[1].getMessage(),
322403
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
323404
)
405+
spans = self.memory_exporter.get_finished_spans()
406+
self.assertEqual(len(spans), 2)
407+
# first span is connection to db
408+
self.assertEqual(spans[0].name, "connect")
409+
# second span is query itself
410+
query_span = spans[1]
411+
self.assertEqual(
412+
query_span.attributes[SpanAttributes.DB_STATEMENT],
413+
"SELECT 1;",
414+
)
415+
416+
asyncio.get_event_loop().run_until_complete(run())
417+
418+
@pytest.mark.skipif(
419+
not sqlalchemy.__version__.startswith("1.4"),
420+
reason="only run async tests for 1.4",
421+
)
422+
def test_create_async_engine_wrapper_enable_commenter_stmt_enabled(self):
423+
async def run():
424+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
425+
SQLAlchemyInstrumentor().instrument(
426+
enable_commenter=True,
427+
commenter_options={
428+
"db_framework": False,
429+
},
430+
enable_attribute_commenter=True,
431+
)
432+
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
433+
create_async_engine,
434+
)
435+
436+
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
437+
async with engine.connect() as cnx:
438+
await cnx.execute(text("SELECT 1;"))
439+
# sqlcommenter
440+
self.assertRegex(
441+
self.caplog.records[1].getMessage(),
442+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
443+
)
444+
spans = self.memory_exporter.get_finished_spans()
445+
self.assertEqual(len(spans), 2)
446+
# first span is connection to db
447+
self.assertEqual(spans[0].name, "connect")
448+
# second span is query itself
449+
query_span = spans[1]
450+
self.assertRegex(
451+
query_span.attributes[SpanAttributes.DB_STATEMENT],
452+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
453+
)
324454

325455
asyncio.get_event_loop().run_until_complete(run())
326456

@@ -352,6 +482,57 @@ async def run():
352482
self.caplog.records[1].getMessage(),
353483
r"SELECT 1 /\*db_driver='(.*)'\*/;",
354484
)
485+
spans = self.memory_exporter.get_finished_spans()
486+
self.assertEqual(len(spans), 2)
487+
# first span is connection to db
488+
self.assertEqual(spans[0].name, "connect")
489+
# second span is query itself
490+
query_span = spans[1]
491+
self.assertEqual(
492+
query_span.attributes[SpanAttributes.DB_STATEMENT],
493+
"SELECT 1;",
494+
)
495+
496+
asyncio.get_event_loop().run_until_complete(run())
497+
498+
@pytest.mark.skipif(
499+
not sqlalchemy.__version__.startswith("1.4"),
500+
reason="only run async tests for 1.4",
501+
)
502+
def test_create_async_engine_wrapper_enable_commenter_stmt_enabled_otel_values_false(
503+
self,
504+
):
505+
async def run():
506+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
507+
SQLAlchemyInstrumentor().instrument(
508+
enable_commenter=True,
509+
commenter_options={
510+
"db_framework": False,
511+
"opentelemetry_values": False,
512+
},
513+
)
514+
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
515+
create_async_engine,
516+
)
517+
518+
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
519+
async with engine.connect() as cnx:
520+
await cnx.execute(text("SELECT 1;"))
521+
# sqlcommenter
522+
self.assertRegex(
523+
self.caplog.records[1].getMessage(),
524+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
525+
)
526+
spans = self.memory_exporter.get_finished_spans()
527+
self.assertEqual(len(spans), 2)
528+
# first span is connection to db
529+
self.assertEqual(spans[0].name, "connect")
530+
# second span is query itself
531+
query_span = spans[1]
532+
self.assertRegex(
533+
query_span.attributes[SpanAttributes.DB_STATEMENT],
534+
r"SELECT 1 /\*db_driver='(.*)'*/;",
535+
)
355536

356537
asyncio.get_event_loop().run_until_complete(run())
357538

0 commit comments

Comments
 (0)