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 30d7b96 commit a05acca
Showing 1 changed file with 173 additions and 1 deletion.
174 changes: 173 additions & 1 deletion instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,59 @@ def test_sqlcommenter_enabled(self):
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

def test_sqlcommenter_enabled_stmt_disabled_default_matches_db_statement_attribute(
def test_sqlcommenter_default_stmt_enabled_no_comments_anywhere(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
# enable_commenter not set
enable_attribute_commenter=True,
)
cnx = engine.connect()
cnx.execute(text("SELECT 1;")).fetchall()
query_log = self.caplog.records[-2].getMessage()
self.assertEqual(
query_log,
"SELECT 1;",
)
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_sqlcommenter_disabled_stmt_enabled_no_comments_anywhere(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
enable_commenter=False,
enable_attribute_commenter=True,
)
cnx = engine.connect()
cnx.execute(text("SELECT 1;")).fetchall()
query_log = self.caplog.records[-2].getMessage()
self.assertEqual(
query_log,
"SELECT 1;",
)
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_sqlcommenter_enabled_stmt_disabled_default(
self,
):
engine = create_engine("sqlite:///:memory:")
Expand Down Expand Up @@ -142,6 +194,45 @@ def test_sqlcommenter_enabled_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],
r"SELECT 1;",
)

def test_sqlcommenter_enabled_stmt_enabled_otel_values_false(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
enable_commenter=True,
commenter_options={
"db_framework": False,
"opentelemetry_values": False,
},
enable_attribute_commenter=True,
)
cnx = engine.connect()
cnx.execute(text("SELECT 1;")).fetchall()
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_sqlcommenter_flask_integration(self):
engine = create_engine("sqlite:///:memory:")
Expand All @@ -164,6 +255,49 @@ def test_sqlcommenter_flask_integration(self):
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',flask=1,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_sqlcommenter_stmt_enabled_flask_integration(self):
engine = create_engine("sqlite:///:memory:")
SQLAlchemyInstrumentor().instrument(
engine=engine,
tracer_provider=self.tracer_provider,
enable_commenter=True,
commenter_options={"db_framework": False},
enable_attribute_commenter=True,
)
cnx = engine.connect()

current_context = context.get_current()
sqlcommenter_context = context.set_value(
"SQLCOMMENTER_ORM_TAGS_AND_VALUES", {"flask": 1}, current_context
)
context.attach(sqlcommenter_context)

cnx.execute(text("SELECT 1;")).fetchall()
self.assertRegex(
self.caplog.records[-2].getMessage(),
r"SELECT 1 /\*db_driver='(.*)',flask=1,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='(.*)',flask=1,traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
)

def test_sqlcommenter_enabled_create_engine_after_instrumentation(self):
SQLAlchemyInstrumentor().instrument(
Expand All @@ -179,6 +313,44 @@ def test_sqlcommenter_enabled_create_engine_after_instrumentation(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_sqlcommenter_enabled_stmt_enabled_create_engine_after_instrumentation(
self,
):
SQLAlchemyInstrumentor().instrument(
tracer_provider=self.tracer_provider,
enable_commenter=True,
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()
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_sqlcommenter_disabled_create_engine_after_instrumentation(self):
SQLAlchemyInstrumentor().instrument(
Expand Down

0 comments on commit a05acca

Please sign in to comment.