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

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

diff --git a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py
index 9a6a5ff2fa..6d38b1c34c 100644
--- a/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py
+++ b/instrumentation/opentelemetry-instrumentation-psycopg2/tests/test_psycopg2_integration.py
@@ -261,6 +261,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.psycopg2.psycopg2.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.apilevel",
+            "123",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.paramstyle",
+            "test",
+        ):
+            cnx = psycopg2.connect(database="test")
+            cnx = Psycopg2Instrumentor().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='psycopg2%%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.psycopg2.psycopg2.__version__",
+            "foobar",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.__libpq_version__",
+            "foobaz",
+        ), mock.patch(
+            "opentelemetry.instrumentation.psycopg2.psycopg2.threadsafety",
+            "123",
+        ):
+            cnx = psycopg2.connect(database="test")
+            cnx = Psycopg2Instrumentor().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='psycopg2%%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 = psycopg2.connect(database="test")
@@ -271,6 +339,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 = psycopg2.connect(database="test")
+        cnx = Psycopg2Instrumentor().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 = psycopg2.connect(database="test")
+        cnx = Psycopg2Instrumentor().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",
+        )
+
     def test_no_op_tracer_provider(self):
         Psycopg2Instrumentor().instrument(
             tracer_provider=trace.NoOpTracerProvider()