Add tests

Author: tammy-baylis-swi

instrumentation/opentelemetry-instrumentation-pymysql/tests/test_pymysql_integration.py

@@ -115,7 +115,7 @@ def test_instrument_connection(self, mock_connect):
     @mock.patch("opentelemetry.instrumentation.dbapi.instrument_connection")
     @mock.patch("pymysql.connect")
     # pylint: disable=unused-argument
-    def test_instrument_connection_enable_commenter(
+    def test_instrument_connection_enable_commenter_dbapi_kwargs(
         self,
         mock_connect,
         mock_instrument_connection,
@@ -132,10 +132,109 @@ def test_instrument_connection_enable_commenter(
         self.assertEqual(kwargs["enable_commenter"], True)
         self.assertEqual(kwargs["commenter_options"], {"foo": True})
 
+    def test_instrument_connection_with_dbapi_sqlcomment_enabled(self):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            cnx_proxy = PyMySQLInstrumentor().instrument_connection(
+                mock_connection,
+                enable_commenter=True,
+            )
+            cnx_proxy.cursor().execute("Select 1;")
+
+            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(
+                mock_cursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
+            )
+
+    def test_instrument_connection_with_dbapi_sqlcomment_enabled_with_options(
+        self,
+    ):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            cnx_proxy = PyMySQLInstrumentor().instrument_connection(
+                mock_connection,
+                enable_commenter=True,
+                commenter_options={
+                    "dbapi_level": False,
+                    "dbapi_threadsafety": True,
+                    "driver_paramstyle": False,
+                },
+            )
+            cnx_proxy.cursor().execute("Select 1;")
+
+            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(
+                mock_cursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_threadsafety='123',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
+            )
+
+    def test_instrument_connection_with_dbapi_sqlcomment_not_enabled_default(
+        self,
+    ):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            cnx_proxy = PyMySQLInstrumentor().instrument_connection(
+                mock_connection,
+            )
+            cnx_proxy.cursor().execute("Select 1;")
+            self.assertEqual(
+                mock_cursor.execute.call_args[0][0],
+                "Select 1;",
+            )
+
     @mock.patch("opentelemetry.instrumentation.dbapi.wrap_connect")
     @mock.patch("pymysql.connect")
     # pylint: disable=unused-argument
-    def test__instrument_enable_commenter(
+    def test_instrument_enable_commenter_dbapi_kwargs(
         self,
         mock_connect,
         mock_wrap_connect,
@@ -148,6 +247,109 @@ def test__instrument_enable_commenter(
         self.assertEqual(kwargs["enable_commenter"], True)
         self.assertEqual(kwargs["commenter_options"], {"foo": True})
 
+    def test_instrument_with_dbapi_sqlcomment_enabled(
+        self,
+    ):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            PyMySQLInstrumentor()._instrument(
+                enable_commenter=True,
+            )
+            cnx = mock_connect_module.connect(database="test")
+            cursor = cnx.cursor()
+            cursor.execute("Select 1;")
+
+            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(
+                mock_cursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_level='123',dbapi_threadsafety='123',driver_paramstyle='test',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
+            )
+
+    def test_instrument_with_dbapi_sqlcomment_enabled_with_options(
+        self,
+    ):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            PyMySQLInstrumentor()._instrument(
+                enable_commenter=True,
+                commenter_options={
+                    "dbapi_level": False,
+                    "dbapi_threadsafety": True,
+                    "driver_paramstyle": False,
+                },
+            )
+            cnx = mock_connect_module.connect(database="test")
+            cursor = cnx.cursor()
+            cursor.execute("Select 1;")
+
+            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(
+                mock_cursor.execute.call_args[0][0],
+                f"Select 1 /*db_driver='pymysql%%3Afoobar',dbapi_threadsafety='123',mysql_client_version='foobaz',traceparent='00-{trace_id}-{span_id}-01'*/;",
+            )
+
+    def test_instrument_with_dbapi_sqlcomment_not_enabled_default(
+        self,
+    ):
+        mock_connect_module = mock.MagicMock(
+            __name__="pymysql",
+            __version__="foobar",
+            threadsafety="123",
+            apilevel="123",
+            paramstyle="test",
+        )
+        mock_connect_module.get_client_info.return_value = "foobaz"
+        mock_cursor = mock_connect_module.connect().cursor()
+        mock_connection = mock.MagicMock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        with mock.patch(
+            "opentelemetry.instrumentation.pymysql.pymysql",
+            mock_connect_module,
+        ):
+            PyMySQLInstrumentor()._instrument()
+            cnx = mock_connect_module.connect(database="test")
+            cursor = cnx.cursor()
+            cursor.execute("Select 1;")
+            self.assertEqual(
+                mock_cursor.execute.call_args[0][0],
+                "Select 1;",
+            )
+
     @mock.patch("pymysql.connect")
     # pylint: disable=unused-argument
     def test_uninstrument_connection(self, mock_connect):