def trace_integration(tracer=None):
"""Wrap the mysql connector to trace it."""
logging.info('Integrated module: {}'.format(MODULE_NAME))
conn_func = getattr(pymysql, CONN_WRAP_METHOD)
conn_module = inspect.getmodule(conn_func)
wrapped = trace.wrap_conn(conn_func)
setattr(conn_module, CONN_WRAP_METHOD, wrapped)
def test_wrap_conn(self):
wrap_func_name = 'wrap func'
cursor_func_name = 'cursor_func'
def mock_wrap(func):
return wrap_func_name
mock_func = mock.Mock()
mock_return = mock.Mock()
mock_cursor_func = mock.Mock()
mock_cursor_func.__name__ = cursor_func_name
mock_return.cursor = mock_cursor_func
mock_func.return_value = mock_return
wrapped = trace.wrap_conn(mock_func)
patch_wrap = mock.patch(
'opencensus.trace.ext.dbapi.trace.wrap_cursor',
side_effect=mock_wrap)
with patch_wrap:
wrapped()
self.assertEqual(wrapped.__class__.__name__, 'function')
self.assertEqual(
getattr(mock_return, cursor_func_name, None), wrap_func_name)