Skip to content

Commit

Permalink
Add tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tammy-baylis-swi
tammy-baylis-swi committed Dec 17, 2024
1 parent a05acca commit d6b625e
Showing 1 changed file with 181 additions and 0 deletions.
181 changes: 181 additions & 0 deletions instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,44 @@ def test_create_engine_wrapper_enable_commenter(self):
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertEqual(
query_span.attributes[SpanAttributes.DB_STATEMENT],
"SELECT 1;",
)

def test_create_engine_wrapper_enable_commenter_stmt_enabled(self):
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={"db_framework": False},
enable_attribute_commenter=True,
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute(text("SELECT 1;")).fetchall()
# sqlcommenter
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertRegex(
query_span.attributes[SpanAttributes.DB_STATEMENT],
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
Expand All @@ -229,6 +267,49 @@ def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertEqual(
query_span.attributes[SpanAttributes.DB_STATEMENT],
"SELECT 1;",
)

def test_create_engine_wrapper_enable_commenter_stmt_enabled_otel_values_false(
self,
):
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
enable_attribute_commenter=True,
)
from sqlalchemy import create_engine # pylint: disable-all

engine = create_engine("sqlite:///:memory:")
cnx = engine.connect()
cnx.execute(text("SELECT 1;")).fetchall()
# sqlcommenter
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertRegex(
query_span.attributes[SpanAttributes.DB_STATEMENT],
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)

def test_custom_tracer_provider(self):
provider = TracerProvider(
Expand Down Expand Up @@ -321,6 +402,55 @@ async def run():
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertEqual(
query_span.attributes[SpanAttributes.DB_STATEMENT],
"SELECT 1;",
)

asyncio.get_event_loop().run_until_complete(run())

@pytest.mark.skipif(
not sqlalchemy.__version__.startswith("1.4"),
reason="only run async tests for 1.4",
)
def test_create_async_engine_wrapper_enable_commenter_stmt_enabled(self):
async def run():
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
},
enable_attribute_commenter=True,
)
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
create_async_engine,
)

engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.connect() as cnx:
await cnx.execute(text("SELECT 1;"))
# sqlcommenter
self.assertRegex(
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertRegex(
query_span.attributes[SpanAttributes.DB_STATEMENT],
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

asyncio.get_event_loop().run_until_complete(run())

Expand Down Expand Up @@ -352,6 +482,57 @@ async def run():
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertEqual(
query_span.attributes[SpanAttributes.DB_STATEMENT],
"SELECT 1;",
)

asyncio.get_event_loop().run_until_complete(run())

@pytest.mark.skipif(
not sqlalchemy.__version__.startswith("1.4"),
reason="only run async tests for 1.4",
)
def test_create_async_engine_wrapper_enable_commenter_stmt_enabled_otel_values_false(
self,
):
async def run():
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
SQLAlchemyInstrumentor().instrument(
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
)
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
create_async_engine,
)

engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.connect() as cnx:
await cnx.execute(text("SELECT 1;"))
# sqlcommenter
self.assertRegex(
self.caplog.records[1].getMessage(),
r"SELECT 1 /\*db_driver='(.*)'\*/;",
)
spans = self.memory_exporter.get_finished_spans()
self.assertEqual(len(spans), 2)
# first span is connection to db
self.assertEqual(spans[0].name, "connect")
# second span is query itself
query_span = spans[1]
self.assertRegex(
query_span.attributes[SpanAttributes.DB_STATEMENT],
r"SELECT 1 /\*db_driver='(.*)'*/;",
)

asyncio.get_event_loop().run_until_complete(run())

Expand Down

0 comments on commit d6b625e

Please sign in to comment.