def test_instrument_connection(self): connection = mock.Mock() # Avoid get_attributes failing because can't concatenate mock connection.database = "-" connection2 = dbapi.instrument_connection(self.tracer, connection, "-") self.assertIsInstance(connection2, dbapi.TracedConnectionProxy) self.assertIs(connection2.__wrapped__, connection)
def test_uninstrument_connection(self): connection = mock.Mock() # Set connection.database to avoid a failure because mock can't # be concatenated connection.database = "-" connection2 = dbapi.instrument_connection(self.tracer, connection, "-") self.assertIs(connection2.__wrapped__, connection) connection3 = dbapi.uninstrument_connection(connection2) self.assertIs(connection3, connection) with self.assertLogs(level=logging.WARNING): connection4 = dbapi.uninstrument_connection(connection) self.assertIs(connection4, connection)
def instrument_connection(self, connection): """Enable instrumentation in a PyMySQL connection. Args: connection: The connection to instrument. Returns: An instrumented connection. """ tracer = get_tracer(__name__, __version__) return dbapi.instrument_connection( tracer, connection, self._DATABASE_COMPONENT, self._DATABASE_TYPE, self._CONNECTION_ATTRIBUTES, )