diff --git a/instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py b/instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py index 6f8af5d8df..18b1f0329b 100644 --- a/instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py +++ b/instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py @@ -111,6 +111,42 @@ def test_instrument_connection(self, mock_connect): spans_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans_list), 1) + @mock.patch("opentelemetry.instrumentation.dbapi.instrument_connection") + @mock.patch("pymysql.connect") + # pylint: disable=unused-argument + def test_instrument_connection_enable_commenter( + self, + mock_connect, + mock_instrument_connection, + ): + cnx = pymysql.connect(database="test") + cnx = PyMySQLInstrumentor().instrument_connection( + cnx, + enable_commenter=True, + commenter_options={"foo": True}, + ) + cursor = cnx.cursor() + cursor.execute("SELECT * FROM test") + kwargs = mock_instrument_connection.call_args[1] + self.assertEqual(kwargs["enable_commenter"], True) + self.assertEqual(kwargs["commenter_options"], {"foo": True}) + + @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect") + @mock.patch("pymysql.connect") + # pylint: disable=unused-argument + def test__instrument_enable_commenter( + self, + mock_connect, + mock_wrap_connect, + ): + PyMySQLInstrumentor()._instrument( + enable_commenter=True, + commenter_options={"foo": True}, + ) + kwargs = mock_wrap_connect.call_args[1] + self.assertEqual(kwargs["enable_commenter"], True) + self.assertEqual(kwargs["commenter_options"], {"foo": True}) + @mock.patch("pymysql.connect") # pylint: disable=unused-argument def test_uninstrument_connection(self, mock_connect):