From 4f969efbf1d62fa02a490dac218b5ff8cd92beb3 Mon Sep 17 00:00:00 2001
From: tammy-baylis-swi <tammy.baylis@solarwinds.com>
Date: Fri, 6 Dec 2024 17:44:58 -0800
Subject: [PATCH] Add psycopg2 tests

---
 .../tests/test_psycopg_integration.py         | 95 +++++++++++++++++++
 1 file changed, 95 insertions(+)

diff --git a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py
index 4ddaad9174..4081070aeb 100644
--- a/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py
+++ b/instrumentation/opentelemetry-instrumentation-psycopg/tests/test_psycopg_integration.py
@@ -378,6 +378,74 @@ def test_sqlcommenter_enabled(self, event_mocked):
         kwargs = event_mocked.call_args[1]
         self.assertEqual(kwargs["enable_commenter"], True)
 
+    def test_sqlcommenter_enabled_instrument_connection_defaults(self):
+        with mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.apilevel",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.paramstyle",
+            "test",
+        ):
+            cnx = psycopg.connect(database="test")
+            cnx = PsycopgInstrumentor().instrument_connection(
+                cnx,
+                enable_commenter=True,
+            )
+            query = "Select 1"
+            cursor = cnx.cursor()
+            cursor.execute(query)
+            spans_list = self.memory_exporter.get_finished_spans()
+            span = spans_list[0]
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            self.assertEqual(
+                MockCursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+
+    def test_sqlcommenter_enabled_instrument_connection_with_options(self):
+        with mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.pq.__build_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg.psycopg.threadsafety",
+            "123",
+        ):
+            cnx = psycopg.connect(database="test")
+            cnx = PsycopgInstrumentor().instrument_connection(
+                cnx,
+                enable_commenter=True,
+                commenter_options={
+                    "dbapi_level": False,
+                    "dbapi_threadsafety": True,
+                    "driver_paramstyle": False,
+                    "foo": "ignored",
+                },
+            )
+            query = "Select 1"
+            cursor = cnx.cursor()
+            cursor.execute(query)
+            spans_list = self.memory_exporter.get_finished_spans()
+            span = spans_list[0]
+            span_id = format(span.get_span_context().span_id, "016x")
+            trace_id = format(span.get_span_context().trace_id, "032x")
+            self.assertEqual(
+                MockCursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='psycopg%%3Afoobar',dbapi_threadsafety='123',libpq_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/",
+            )
+
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     def test_sqlcommenter_disabled(self, event_mocked):
         cnx = psycopg.connect(database="test")
@@ -388,6 +456,33 @@ def test_sqlcommenter_disabled(self, event_mocked):
         kwargs = event_mocked.call_args[1]
         self.assertEqual(kwargs["enable_commenter"], False)
 
+    def test_sqlcommenter_disabled_default_instrument_connection(self):
+        cnx = psycopg.connect(database="test")
+        cnx = PsycopgInstrumentor().instrument_connection(
+            cnx,
+        )
+        query = "Select 1"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+        self.assertEqual(
+            MockCursor.execute.call_args[0][0],
+            "Select 1",
+        )
+
+    def test_sqlcommenter_disabled_explicit_instrument_connection(self):
+        cnx = psycopg.connect(database="test")
+        cnx = PsycopgInstrumentor().instrument_connection(
+            cnx,
+            enable_commenter=False,
+        )
+        query = "Select 1"
+        cursor = cnx.cursor()
+        cursor.execute(query)
+        self.assertEqual(
+            MockCursor.execute.call_args[0][0],
+            "Select 1",
+        )
+
 
 class TestPostgresqlIntegrationAsync(
     PostgresqlIntegrationTestMixin, TestBase, IsolatedAsyncioTestCase