def patch(): if getattr(elasticsearch, '_datadog_patch', False): return setattr(elasticsearch, '_datadog_patch', True) wrapt.wrap_function_wrapper('elasticsearch.transport', 'Transport.perform_request', _perform_request) Pin(service=DEFAULT_SERVICE, app="elasticsearch", app_type="db").onto(elasticsearch.transport.Transport)
def patch(): if getattr(botocore.client, '_datadog_patch', False): return setattr(botocore.client, '_datadog_patch', True) wrapt.wrap_function_wrapper('botocore.client', 'BaseClient._make_api_call', patched_api_call) Pin(service="aws", app="aws", app_type="web").onto(botocore.client.BaseClient)
def _patch_extensions(_extensions): # we must patch extensions all the time (it's pretty harmless) so split # from global patching of connections. must be idempotent. for _, module, func, wrapper in _extensions: if not hasattr(module, func) or isinstance(getattr(module, func), wrapt.ObjectProxy): continue wrapt.wrap_function_wrapper(module, func, wrapper)
def _instrument( tracer_provider: TracerProvider = None, url_filter: _UrlFilterT = None, span_name: _SpanNameT = None, ): """Enables tracing of all ClientSessions When a ClientSession gets created a TraceConfig is automatically added to the session's trace_configs. """ # pylint:disable=unused-argument def instrumented_init(wrapped, instance, args, kwargs): if context_api.get_value("suppress_instrumentation"): return wrapped(*args, **kwargs) trace_configs = list(kwargs.get("trace_configs") or ()) trace_config = create_trace_config( url_filter=url_filter, span_name=span_name, tracer_provider=tracer_provider, ) trace_config.opentelemetry_aiohttp_instrumented = True trace_configs.append(trace_config) kwargs["trace_configs"] = trace_configs return wrapped(*args, **kwargs) wrapt.wrap_function_wrapper( aiohttp.ClientSession, "__init__", instrumented_init )
def trace_pyramid(config): config.add_tween('ddtrace.contrib.pyramid:trace_tween_factory') # ensure we only patch the renderer once. if not isinstance(pyramid.renderers.RendererHelper.render, wrapt.ObjectProxy): wrapt.wrap_function_wrapper('pyramid.renderers', 'RendererHelper.render', trace_render)
def patch(): """ Patch module. :return: None """ wrapt.wrap_function_wrapper('greengrasssdk.IoTDataPlane', 'Client.publish', _wrapper)
def _instrument( tracer_provider: TracerProvider = None, url_filter: _UrlFilterT = None, request_hook: _RequestHookT = None, response_hook: _ResponseHookT = None, ): """Enables tracing of all ClientSessions When a ClientSession gets created a TraceConfig is automatically added to the session's trace_configs. """ # pylint:disable=unused-argument def instrumented_init(wrapped, instance, args, kwargs): if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY): return wrapped(*args, **kwargs) trace_configs = list(kwargs.get("trace_configs") or ()) trace_config = create_trace_config( url_filter=url_filter, request_hook=request_hook, response_hook=response_hook, tracer_provider=tracer_provider, ) trace_config._is_instrumented_by_opentelemetry = True trace_configs.append(trace_config) kwargs["trace_configs"] = trace_configs return wrapped(*args, **kwargs) wrapt.wrap_function_wrapper(aiohttp.ClientSession, "__init__", instrumented_init)
def hook(module): """ Hook method to install the Instana middleware into Flask """ if "INSTANA_DEV" in os.environ: print("==============================================================") print("Instana: Running flask hook") print("==============================================================") wrapt.wrap_function_wrapper('flask', 'Flask.__init__', wrapper)
def instrument(tracer=None): flask = utils.get_module('flask') if utils.is_instrumented(flask): return flask_opentracing = utils.get_module('flask_opentracing') def flask_tracer(__init__, app, args, kwargs): """ A function wrapper of flask.Flask.__init__ to create a corresponding flask_opentracing.FlaskTracer upon app instantiation. """ __init__(*args, **kwargs) _tracer = tracer or config.tracer or opentracing.tracer app.config['FLASK_TRACER'] = flask_opentracing.FlaskTracer( tracer=_tracer, trace_all_requests=config.trace_all, app=app, traced_attributes=config.traced_attributes) wrap_function_wrapper('flask', 'Flask.__init__', flask_tracer) utils.mark_instrumented(flask)
def patch_method_with_caching(name: str, cls: BoundMethodClassType, save_state: bool) -> None: wrapt.wrap_function_wrapper( module=cls.__module__, name=f"{cls.__name__}.{name}", wrapper=partial(_caching_method_wrapper, save_state=save_state), )
def instrument(tracer=None): """ Requests auto-instrumentation works by hooking a __new__ proxy for a SessionTracing instance upon requests.sessions.Session initialization to trigger proper inheritance. SessionTracing.__init__ is also wrapped for correct argument injection. """ requests = utils.get_module('requests') if utils.is_instrumented(requests): return def session_tracing_init(__init__, instance, _, __): _tracer = tracer or config.tracer or opentracing.tracer __init__(_tracer, propagate=config.propagate, span_tags=config.span_tags or {}) from requests_opentracing import SessionTracing _session_new[0] = requests.Session.__new__ _session_tracing_new[0] = SessionTracing.__new__ SessionTracing.__new__ = session_tracing_new.__get__(SessionTracing) requests.Session.__new__ = session_new.__get__(requests.Session) wrap_function_wrapper('requests_opentracing.tracing', 'SessionTracing.__init__', session_tracing_init) utils.mark_instrumented(requests)
def patch(): """Instrument Pylons applications""" if getattr(pylons.wsgiapp, '_datadog_patch', False): return setattr(pylons.wsgiapp, '_datadog_patch', True) wrapt.wrap_function_wrapper('pylons.wsgiapp', 'PylonsApp.__init__', traced_init)
def test_requests__patch_with_sentry__sunset(http, mocker, sunset_date, sunset_link, expected_message): m_capture_message = mocker.patch("kw.platform.utils.capture_message") url = "http://kiwi.com" headers = {} if sunset_date: headers["Sunset"] = sunset_date if sunset_link: headers["Link"] = '<{}>;rel="sunset"'.format(sunset_link) http.register_uri(http.GET, url, body="Hello", adding_headers=headers) session = requests.Session() session.get(url) m_capture_message.assert_not_called() wrapt.wrap_function_wrapper(session, "request", wrappers.add_sentry_handler()) session.get(url) if expected_message: m_capture_message.assert_called_once_with(expected_message.format( sd=sunset_date, sl=sunset_link), level="warning") else: m_capture_message.assert_not_called()
def patch(): """ Patch module. :return: None """ wrapt.wrap_function_wrapper('aiohttp.web', 'Application.__init__', _wrapper)
def instrument(tracer=None): pymongo = utils.get_module('pymongo') if utils.is_instrumented(pymongo): return pymongo_opentracing = utils.get_module('pymongo_opentracing') def pymongo_tracer(__init__, app, args, kwargs): """ A function wrapper of pymongo.MongoClient.__init__ to register a corresponding pymongo_opentracing.CommandTracing upon client instantiation. """ _tracer = tracer or config.tracer or opentracing.tracer command_tracing = pymongo_opentracing.CommandTracing( tracer=_tracer, span_tags=config.span_tags or {}, ) event_listeners = list(kwargs.pop('event_listeners', [])) event_listeners.insert(0, command_tracing) kwargs['event_listeners'] = event_listeners __init__(*args, **kwargs) wrap_function_wrapper('pymongo', 'MongoClient.__init__', pymongo_tracer) utils.mark_instrumented(pymongo)
def _wrap_send_message(self, sqs_class: type) -> None: def send_wrapper(wrapped, instance, args, kwargs): if context.get_value(_SUPPRESS_INSTRUMENTATION_KEY): return wrapped(*args, **kwargs) queue_url = kwargs.get("QueueUrl") # The method expect QueueUrl and Entries params, so if they are None, we call wrapped to receive the # original exception queue_name = Boto3SQSInstrumentor._extract_queue_name_from_url( queue_url) with self._tracer.start_as_current_span( name=f"{queue_name} send", kind=SpanKind.PRODUCER, end_on_exit=True, ) as span: Boto3SQSInstrumentor._enrich_span(span, queue_name, queue_url) attributes = kwargs.pop("MessageAttributes", {}) propagate.inject(attributes, setter=boto3sqs_setter) retval = wrapped(*args, MessageAttributes=attributes, **kwargs) message_id = retval.get("MessageId") if message_id: if span.is_recording(): span.set_attribute(SpanAttributes.MESSAGING_MESSAGE_ID, message_id) return retval wrap_function_wrapper(sqs_class, "send_message", send_wrapper)
async def test_aiohttp__patch_with_sentry__deprecated_usage( loop, aiomock, mocker, message): m_capture_message = mocker.patch("kw.platform.utils.capture_message") url = "http://kiwi.com" headers = {} if message: headers["Deprecated-Usage"] = message aiomock.get(url, headers=headers) async with aiohttp.ClientSession() as client: async with client.get(url) as resp: await resp.text() m_capture_message.assert_not_called() aiomock.get(url, headers=headers) async with aiohttp.ClientSession() as client: wrapt.wrap_function_wrapper(client, "_request", uut.monkey._add_sentry_handler()) async with client.get(url) as resp: await resp.text() if message: m_capture_message.assert_called_once_with(message, level="warning") else: m_capture_message.assert_not_called()
def __init__(self, plugin_context=None, config=None): self.hooks = { 'before:invocation': self.before_invocation, 'after:invocation': self.after_invocation } self.plugin_context = plugin_context self.logger = logging.getLogger('STDOUT') if isinstance(config, LogConfig): self.config = config else: self.config = LogConfig() with LogPlugin.lock: if (not LogPlugin.wrapped) and (not ConfigProvider.get( config_names.THUNDRA_LOG_CONSOLE_DISABLE)): if PY37 or PY38: wrapt.wrap_function_wrapper('builtins', 'print', self._wrapper) else: sys.stdout = StreamToLogger(self.logger, sys.stdout) LogPlugin.wrapped = True if not ConfigProvider.get(config_names.THUNDRA_LOG_CONSOLE_DISABLE): handler = ThundraLogHandler() has_thundra_log_handler = False for log_handlers in self.logger.handlers: if isinstance(log_handlers, ThundraLogHandler): has_thundra_log_handler = True if not has_thundra_log_handler: self.logger.addHandler(handler) self.logger.setLevel(logging.INFO) handler.setLevel(logging.INFO) self.logger.propagate = False
def patch_pymongo(context): """ Monkey patches pymongo client, if available. Overloads the query methods to add tracing and metrics collection. """ def wrapper(wrapped, instance, args, kwargs): if not hasattr(context, "iopipe") or not hasattr( context.iopipe, "mark"): # pragma: no cover return wrapped(*args, **kwargs) id = ensure_utf8(str(uuid.uuid4())) with context.iopipe.mark(id): response = wrapped(*args, **kwargs) trace = context.iopipe.mark.measure(id) context.iopipe.mark.delete(id) collect_pymongo_metrics(context, trace, instance, response) return response for module, attr, _wrapper in [ ("pymongo.collection", "Collection.find", wrapper), ("pymongo.collection", "Collection.bulk_write", wrapper), ("pymongo.collection", "Collection.delete_many", wrapper), ("pymongo.collection", "Collection.delete_one", wrapper), ("pymongo.collection", "Collection.insert_many", wrapper), ("pymongo.collection", "Collection.insert_one", wrapper), ("pymongo.collection", "Collection.replace_one", wrapper), ("pymongo.collection", "Collection.update_many", wrapper), ("pymongo.collection", "Collection.update_one", wrapper), ]: try: wrapt.wrap_function_wrapper(module, attr, _wrapper) except Exception: # pragma: no cover pass
async def test_aiohttp__patch_with_sentry__sunset(loop, aiomock, mocker, sunset_date, sunset_link, expected_message): m_capture_message = mocker.patch("kw.platform.utils.capture_message") url = "http://kiwi.com" headers = {} if sunset_date: headers["Sunset"] = sunset_date if sunset_link: headers["Link"] = '<{}>;rel="sunset"'.format(sunset_link) aiomock.get(url, headers=headers) async with aiohttp.ClientSession() as client: async with client.get(url) as resp: await resp.text() m_capture_message.assert_not_called() aiomock.get(url, headers=headers) async with aiohttp.ClientSession() as client: wrapt.wrap_function_wrapper(client, "_request", uut.monkey._add_sentry_handler()) async with client.get(url) as resp: await resp.text() if expected_message: m_capture_message.assert_called_once_with(expected_message.format( sd=sunset_date, sl=sunset_link), level="warning") else: m_capture_message.assert_not_called()
def patch_pymysql(context): """ Monkey patches pymysql client, if available. Overloads the execute method to add tracing and metrics collection. """ class _CursorProxy(CursorProxy): def execute(self, *args, **kwargs): if not hasattr(context, "iopipe") or not hasattr( context.iopipe, "mark"): # pragma: no cover self.__wrapped__.execute(*args, **kwargs) return id = ensure_utf8(str(uuid.uuid4())) with context.iopipe.mark(id): self.__wrapped__.execute(*args, **kwargs) trace = context.iopipe.mark.measure(id) context.iopipe.mark.delete(id) collect_mysql_metrics(context, trace, self, args) class _ConnectionProxy(ConnectionProxy): def cursor(self, *args, **kwargs): cursor = self.__wrapped__.cursor(*args, **kwargs) return _CursorProxy(cursor, self) def connect_wrapper(wrapped, instance, args, kwargs): connection = wrapped(*args, **kwargs) return _ConnectionProxy(connection, args, kwargs) try: wrapt.wrap_function_wrapper("pymysql", "connect", connect_wrapper) except Exception: # pragma: no cover pass
def patch(): if getattr(aiobotocore.client, '_datadog_patch', False): return setattr(aiobotocore.client, '_datadog_patch', True) wrapt.wrap_function_wrapper('aiobotocore.client', 'AioBaseClient._make_api_call', _wrapped_api_call) Pin(service='aws', app='aws', app_type='web').onto(aiobotocore.client.AioBaseClient)
def patch(): """ Patch module. :return: None """ wrapt.wrap_function_wrapper('requests', 'Session.send', _wrapper)
def wrap_create_pool( name: str, database_component: str, database_type: str = "", connection_attributes: typing.Dict = None, version: str = "", tracer_provider: typing.Optional[TracerProvider] = None, ): # pylint: disable=unused-argument async def wrap_create_pool_( wrapped: typing.Callable[..., typing.Any], instance: typing.Any, args: typing.Tuple[typing.Any, typing.Any], kwargs: typing.Dict[typing.Any, typing.Any], ): db_integration = AiopgIntegration( name, database_component, database_type, connection_attributes=connection_attributes, version=version, tracer_provider=tracer_provider, ) return await db_integration.wrapped_pool(wrapped, args, kwargs) try: wrapt.wrap_function_wrapper(aiopg, "create_pool", wrap_create_pool_) except Exception as ex: # pylint: disable=broad-except logger.warning("Failed to integrate with DB API. %s", str(ex))
def _instrument(self, **kwargs): self._original_kafka_producer = confluent_kafka.Producer self._original_kafka_consumer = confluent_kafka.Consumer confluent_kafka.Producer = AutoInstrumentedProducer confluent_kafka.Consumer = AutoInstrumentedConsumer tracer_provider = kwargs.get("tracer_provider") tracer = trace.get_tracer( __name__, __version__, tracer_provider=tracer_provider ) self._tracer = tracer def _inner_wrap_produce(func, instance, args, kwargs): return ConfluentKafkaInstrumentor.wrap_produce( func, instance, self._tracer, args, kwargs ) def _inner_wrap_poll(func, instance, args, kwargs): return ConfluentKafkaInstrumentor.wrap_poll( func, instance, self._tracer, args, kwargs ) wrapt.wrap_function_wrapper( AutoInstrumentedProducer, "produce", _inner_wrap_produce, ) wrapt.wrap_function_wrapper( AutoInstrumentedConsumer, "poll", _inner_wrap_poll, )
def monkey_patch_botocore_for_xray(): """Explicit way to monkey-patch botocore to trace AWS API calls.""" wrapt.wrap_function_wrapper( 'botocore.client', 'BaseClient._make_api_call', xray_botocore_api_call, )
def patch(): wrapt.wrap_function_wrapper('requests', 'Session.request', _xray_traced_requests) wrapt.wrap_function_wrapper('requests', 'Session.prepare_request', _inject_header)
def monkey_patch_requests_for_xray(): """Explicit way to monkey-patch requests to trace HTTP requests.""" wrapt.wrap_function_wrapper( 'requests.sessions', 'Session.send', xray_requests_send, )
def patch() -> None: """ This patches the :code:`Service._dispatch_future` and :code:`Service._dispatch_send` methods. The :code:`dispatch_future` method is patched to extract the :code:`skip_breaker` keyword argument because current insanic doesn't pass the :code:`kwargs` to :code:`_dispatch_send`. Maybe if this is fixed in Insanic, this can be removed" """ if not hasattr(Service._dispatch_future, "__wrapped__"): wrapt.wrap_function_wrapper( "insanic.services", "Service._dispatch_future", request_breaker.extract_skip_breaker, ) wrapt.wrap_function_wrapper( "insanic.services", "Service._dispatch_send", request_breaker.wrapped_request, )
def instrument(tracer=None): tornado = utils.get_module('tornado') if utils.is_instrumented(tornado): return tornado_opentracing = utils.get_module('tornado_opentracing') def _tracer_config(wrapped_tracer_config, _, wrapt_args, __): """ A function wrapper for tornado_opentracing's monkey patcher of tornado.web.Application.__init__() used to inject tracer configuration as settings arguments. As a wrapt function wrapper of a function_wrapper, _tracer_config's meaningful arguments are oddly nested. """ __init__ = wrapt_args[0] app = wrapt_args[1] args = wrapt_args[2] kwargs = wrapt_args[3] _tracer = tracer or config.tracer or opentracing.tracer kwargs['opentracing_tracing'] = tornado_opentracing.TornadoTracing(_tracer) kwargs['opentracing_trace_all'] = config.trace_all kwargs['opentracing_trace_client'] = config.trace_client kwargs['opentracing_traced_attributes'] = config.traced_attributes kwargs['opentracing_start_span_cb'] = config.start_span_cb wrapped_tracer_config(__init__, app, args, kwargs) wrap_function_wrapper('tornado_opentracing.application', 'tracer_config', _tracer_config) tornado_opentracing.init_tracing() utils.mark_instrumented(tornado)
def instrument(tracer=None): """ Requests auto-instrumentation works by hooking a __new__ proxy for a CeleryTracing instance upon celery.Celery initialization to trigger proper inheritance. CeleryTracing.__init__ is also wrapped for correct argument injection. """ celery = utils.get_module('celery') if utils.is_instrumented(celery): return import celery.app def celery_tracing_init(__init__, instance, args, kwargs): _tracer = tracer or config.tracer or opentracing.tracer __init__(*args, tracer=_tracer, propagate=config.propagate, span_tags=config.span_tags or {}, **kwargs) from celery_opentracing import CeleryTracing _celery_new[0] = celery.Celery.__new__ _celery_tracing_new[0] = CeleryTracing.__new__ CeleryTracing.__new__ = celery_tracing_new.__get__(CeleryTracing) celery.app.base.Celery.__new__ = celery_new.__get__(celery.Celery) wrap_function_wrapper('celery_opentracing.tracing', 'CeleryTracing.__init__', celery_tracing_init) utils.mark_instrumented(celery)
def _instrument( tracer, span_name_or_callback: _SpanNameT = None, url_filter: _UrlFilterT = None, ): def instrumented_urlopen(wrapped, instance, args, kwargs): if _is_instrumentation_suppressed(): return wrapped(*args, **kwargs) method = _get_url_open_arg("method", args, kwargs).upper() url = _get_url(instance, args, kwargs, url_filter) headers = _prepare_headers(kwargs) span_name = _get_span_name(span_name_or_callback, method, url, headers) span_attributes = { SpanAttributes.HTTP_METHOD: method, SpanAttributes.HTTP_URL: url, } with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT, attributes=span_attributes) as span: inject(headers) with _suppress_further_instrumentation(): response = wrapped(*args, **kwargs) _apply_response(span, response) return response wrapt.wrap_function_wrapper( urllib3.connectionpool.HTTPConnectionPool, "urlopen", instrumented_urlopen, )
def patch(): """Patch the instrumented Flask object """ if getattr(flask, '_datadog_patch', False): return setattr(flask, '_datadog_patch', True) wrapt.wrap_function_wrapper('flask', 'Flask.__init__', traced_init)
def patch(): """Patch the bottle.Bottle class """ if getattr(bottle, '_datadog_patch', False): return setattr(bottle, '_datadog_patch', True) wrapt.wrap_function_wrapper('bottle', 'Bottle.__init__', traced_init)
def patch(): """ Patch monkey patches psycopg's connection function so that the connection's functions are traced. """ if getattr(psycopg2, '_datadog_patch', False): return setattr(psycopg2, '_datadog_patch', True) wrapt.wrap_function_wrapper(psycopg2, 'connect', patched_connect) _patch_extensions(_psycopg2_extensions) # do this early just in case
def patch(): """ AWSQueryConnection and AWSAuthConnection are two different classes called by different services for connection. For exemple EC2 uses AWSQueryConnection and S3 uses AWSAuthConnection """ if getattr(boto.connection, '_datadog_patch', False): return setattr(boto.connection, '_datadog_patch', True) wrapt.wrap_function_wrapper('boto.connection', 'AWSQueryConnection.make_request', patched_query_request) wrapt.wrap_function_wrapper('boto.connection', 'AWSAuthConnection.make_request', patched_auth_request) Pin(service="aws", app="aws", app_type="web").onto(boto.connection.AWSQueryConnection) Pin(service="aws", app="aws", app_type="web").onto(boto.connection.AWSAuthConnection)
def test_wrap_function_module_name(self): _args = (1, 2) _kwargs = {'one': 1, 'two': 2} called = [] def wrapper(wrapped, instance, args, kwargs): called.append((args, kwargs)) self.assertEqual(instance, None) self.assertEqual(args, _args) self.assertEqual(kwargs, _kwargs) return wrapped(*args, **kwargs) wrapt.wrap_function_wrapper(__name__, 'global_function_1', wrapper) result = global_function_1(*_args, **_kwargs) self.assertEqual(result, (_args, _kwargs)) self.assertEqual(called[0], (_args, _kwargs))
def test_wrap_class_method_module_name(self): _args = (1, 2) _kwargs = {'one': 1, 'two': 2} called = [] def wrapper(wrapped, instance, args, kwargs): called.append((args, kwargs)) self.assertEqual(instance, Class_2) self.assertEqual(args, _args) self.assertEqual(kwargs, _kwargs) return wrapped(*args, **kwargs) wrapt.wrap_function_wrapper(__name__, 'Class_2.method', wrapper) result = Class_2.method(*_args, **_kwargs) self.assertEqual(result, (Class_2, _args, _kwargs)) self.assertEqual(called[0], (_args, _kwargs))
"--os-profile": ("SECRET_KEY", True, False), "--os-interface": (DEFAULT_INTERFACE, True, True), } # Wrap the osc_lib make_shell() function to set the shell class since # osc-lib's TestShell class doesn't allow us to specify it yet. # TODO(dtroyer): remove this once the shell_class_patch patch is released # in osc-lib def make_shell_wrapper(func, inst, args, kwargs): if "shell_class" not in kwargs: kwargs["shell_class"] = shell.OpenStackShell return func(*args, **kwargs) wrapt.wrap_function_wrapper(osc_lib_test_utils, "make_shell", make_shell_wrapper) class TestShell(osc_lib_test_utils.TestShell): # Full name of the OpenStackShell class to test (cliff.app.App subclass) shell_class_name = "openstackclient.shell.OpenStackShell" # TODO(dtroyer): remove this once the shell_class_patch patch is released # in osc-lib app_patch = shell_class_name def setUp(self): super(TestShell, self).setUp() # TODO(dtroyer): remove this once the shell_class_patch patch is # released in osc-lib
def patch(): wrapt.wrap_function_wrapper('pymysql', 'connect', _connect)
def includeme(config): # Add our tween just before the default exception handler config.add_tween(DD_TWEEN_NAME, over=pyramid.tweens.EXCVIEW) # ensure we only patch the renderer once. if not isinstance(pyramid.renderers.RendererHelper.render, wrapt.ObjectProxy): wrapt.wrap_function_wrapper('pyramid.renderers', 'RendererHelper.render', trace_render)
} # Wrap the osc_lib make_shell() function to set the shell class since # osc-lib's TestShell class doesn't allow us to specify it yet. # TODO(dtroyer): remove this once the shell_class_patch patch is released # in osc-lib def make_shell_wrapper(func, inst, args, kwargs): if 'shell_class' not in kwargs: kwargs['shell_class'] = shell.OpenStackShell return func(*args, **kwargs) wrapt.wrap_function_wrapper( osc_lib_test_utils, 'make_shell', make_shell_wrapper, ) class TestShell(osc_lib_test_utils.TestShell): # Full name of the OpenStackShell class to test (cliff.app.App subclass) shell_class_name = "openstackclient.shell.OpenStackShell" # TODO(dtroyer): remove this once the shell_class_patch patch is released # in osc-lib app_patch = shell_class_name def setUp(self): super(TestShell, self).setUp()
def patch(): wrapt.wrap_function_wrapper('mysql.connector', 'connect', _connect) # `Connect` is an alias for `connect`, patch it too if hasattr(mysql.connector, 'Connect'): mysql.connector.Connect = mysql.connector.connect