Integrate only when the object has been initialized

Previous integrations in the SDK were done by importing rollbar.contrib.'framework'.

Since `rollbar._build_person_data()` needs to import Starlette integration,
this sets the Starlette hook for the SDK even if not explicility integrated.

Under certain cicrumstances, this can result in erroneously add Starlette
framework information to the payload and associate occurence with Starlette
integration.

For now, integration with SDK is make only if integration function/class is called.

Fix [ch85137]

rollbar/contrib/asgi/__init__.py

@@ -1,12 +1,3 @@
 __all__ = ['ReporterMiddleware']
 
-import rollbar
-
 from .middleware import ReporterMiddleware
-
-
-def _hook(request, data):
-    data['framework'] = 'asgi'
-
-
-rollbar.BASE_DATA_HOOK = _hook

rollbar/contrib/asgi/middleware.py

@@ -1,12 +1,16 @@
 import sys
 
 import rollbar
+from .integration import IntegrationBase, integrate
 from .types import ASGIApp, Receive, Scope, Send
 from rollbar.lib._async import RollbarAsyncError, try_report
 
 
-class ReporterMiddleware:
+@integrate(framework_name='asgi')
+class ReporterMiddleware(IntegrationBase):
     def __init__(self, app: ASGIApp) -> None:
+        super().__init__()
+
         self.app = app
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

rollbar/contrib/fastapi/__init__.py

@@ -1,18 +1,8 @@
 __all__ = ['add_to', 'ReporterMiddleware', 'LoggerMiddleware', 'get_current_request']
 
-from fastapi import __version__
-
-import rollbar
 from .middleware import ReporterMiddleware
 from .logger import LoggerMiddleware
 from .routing import add_to
 
 # Do not modify the returned request object
 from rollbar.contrib.starlette import get_current_request
-
-
-def _hook(request, data):
-    data['framework'] = f'fastapi {__version__}'
-
-
-rollbar.BASE_DATA_HOOK = _hook

rollbar/contrib/fastapi/logger.py

@@ -1,11 +1,14 @@
 __all__ = ['LoggerMiddleware']
 
 import logging
+from fastapi import __version__
 
+from rollbar.contrib.asgi.integration import integrate
 from rollbar.contrib.starlette import LoggerMiddleware as StarletteLoggerMiddleware
 
 log = logging.getLogger(__name__)
 
 
+@integrate(framework_name=f'fastapi {__version__}')
 class LoggerMiddleware(StarletteLoggerMiddleware):
     ...

rollbar/contrib/fastapi/middleware.py

@@ -1,7 +1,11 @@
 __all__ = ['ReporterMiddleware']
 
+from fastapi import __version__
+
+from rollbar.contrib.asgi.integration import integrate
 from rollbar.contrib.starlette import ReporterMiddleware as StarletteReporterMiddleware
 
 
+@integrate(framework_name=f'fastapi {__version__}')
 class ReporterMiddleware(StarletteReporterMiddleware):
     ...

rollbar/contrib/fastapi/routing.py

@@ -4,7 +4,7 @@
 import sys
 from typing import Callable, Optional, Type, Union
 
-from fastapi import APIRouter, FastAPI
+from fastapi import APIRouter, FastAPI, __version__
 from fastapi.routing import APIRoute
 
 try:
@@ -16,13 +16,15 @@
 
 import rollbar
 from .utils import fastapi_min_version, get_installed_middlewares, has_bare_routing
+from rollbar.contrib.asgi.integration import integrate
 from rollbar.contrib.starlette.requests import store_current_request
 from rollbar.lib._async import RollbarAsyncError, try_report
 
 log = logging.getLogger(__name__)
 
 
 @fastapi_min_version('0.41.0')
+@integrate(framework_name=f'fastapi {__version__}')
 def add_to(app_or_router: Union[FastAPI, APIRouter]) -> Optional[Type[APIRoute]]:
     """
     Adds RollbarLoggingRoute handler to the router app.

rollbar/contrib/starlette/__init__.py

@@ -1,17 +1,7 @@
 __all__ = ['ReporterMiddleware', 'LoggerMiddleware', 'get_current_request']
 
-from starlette import __version__
-
-import rollbar
 from .middleware import ReporterMiddleware
 from .logger import LoggerMiddleware
 
 # Do not modify the returned request object
 from .requests import get_current_request
-
-
-def _hook(request, data):
-    data['framework'] = f'starlette {__version__}'
-
-
-rollbar.BASE_DATA_HOOK = _hook

rollbar/contrib/starlette/logger.py

@@ -3,14 +3,17 @@
 import logging
 import sys
 
+from starlette import __version__
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 from rollbar.contrib.asgi import ReporterMiddleware as ASGIReporterMiddleware
+from rollbar.contrib.asgi.integration import integrate
 from rollbar.contrib.starlette.requests import store_current_request
 
 log = logging.getLogger(__name__)
 
 
+@integrate(framework_name=f'starlette {__version__}')
 class LoggerMiddleware(ASGIReporterMiddleware):
     def __init__(self, app: ASGIApp) -> None:
         if sys.version_info < (3, 7):

rollbar/contrib/starlette/middleware.py

@@ -1,14 +1,17 @@
 import sys
 
+from starlette import __version__
 from starlette.requests import Request
 from starlette.types import Receive, Scope, Send
 
 import rollbar
 from .requests import store_current_request
 from rollbar.contrib.asgi import ReporterMiddleware as ASGIReporterMiddleware
+from rollbar.contrib.asgi.integration import integrate
 from rollbar.lib._async import RollbarAsyncError, try_report
 
 
+@integrate(framework_name=f'starlette {__version__}')
 class ReporterMiddleware(ASGIReporterMiddleware):
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         try:

rollbar/test/asgi_tests/test_integration.py

@@ -1,24 +1,32 @@
-import importlib
-import sys
+from rollbar.test import BaseTest
 
-try:
-    from unittest import mock
-except ImportError:
-    import mock
 
-import unittest2
+class IntegrationTest(BaseTest):
+    def test_should_integrate_if__integrate_defined(self):
+        from rollbar.contrib.asgi.integration import IntegrationBase
 
-import rollbar
-import rollbar.contrib.asgi
-from rollbar.test import BaseTest
+        called = False  # cannot patch local objects
+
+        class Integration(IntegrationBase):
+            def _integrate(self):
+                nonlocal called
+                called = True
+
+        self.assertTrue(hasattr(Integration, '_integrate'))
+
+        obj = Integration()
+
+        self.assertTrue(called)
+        self.assertTrue(hasattr(obj, '_integrate'))
+
+    def test_should_not_fail_if__integrate_not_exists(self):
+        from rollbar.contrib.asgi.integration import IntegrationBase
 
-ALLOWED_PYTHON_VERSION = sys.version_info >= (3, 5)
+        class WrongIntegration(IntegrationBase):
+            ...
 
+        self.assertFalse(hasattr(WrongIntegration, '_integrate'))
 
-@unittest2.skipUnless(ALLOWED_PYTHON_VERSION, 'ASGI implementation requires Python3.5+')
-class ASGIIntegrationTest(BaseTest):
-    def setUp(self):
-        importlib.reload(rollbar.contrib.asgi)
+        obj = WrongIntegration()
 
-    def test_should_set_asgi_hook(self):
-        self.assertEqual(rollbar.BASE_DATA_HOOK, rollbar.contrib.asgi._hook)
+        self.assertFalse(hasattr(obj, '_integrate'))

rollbar/test/asgi_tests/test_middleware.py

@@ -10,7 +10,6 @@
 import unittest2
 
 import rollbar
-import rollbar.contrib.asgi
 from rollbar.lib._async import AsyncMock
 from rollbar.test import BaseTest
 
@@ -23,9 +22,9 @@ class ReporterMiddlewareTest(BaseTest):
     default_settings = copy.deepcopy(rollbar.SETTINGS)
 
     def setUp(self):
+        importlib.reload(rollbar)
         rollbar.SETTINGS = copy.deepcopy(self.default_settings)
         rollbar.SETTINGS['handler'] = 'async'
-        importlib.reload(rollbar.contrib.asgi)
 
     @mock.patch('rollbar.report_exc_info')
     def test_should_catch_and_report_errors(self, mock_report):
@@ -47,6 +46,22 @@ def test_should_catch_and_report_errors(self, mock_report):
         self.assertEqual(exc_type, RuntimeError)
         self.assertIsInstance(exc_value, RuntimeError)
 
+    @mock.patch('rollbar._check_config', return_value=True)
+    @mock.patch('rollbar.send_payload')
+    def test_should_add_framework_name_to_payload(self, mock_send_payload, *mocks):
+        import rollbar
+        from rollbar.contrib.asgi.middleware import ReporterMiddleware
+
+        self.assertIsNone(rollbar.BASE_DATA_HOOK)
+
+        ReporterMiddleware(None)  # invoke integration
+        rollbar.report_exc_info()
+
+        mock_send_payload.assert_called_once()
+        payload = mock_send_payload.call_args[0][0]
+
+        self.assertIn('asgi', payload['data']['framework'])
+
     @unittest2.skipUnless(ASYNC_REPORT_ENABLED, 'Requires Python 3.6+')
     @mock.patch('rollbar.lib._async.report_exc_info', new_callable=AsyncMock)
     @mock.patch('rollbar.report_exc_info')

rollbar/test/fastapi_tests/test_logger.py

@@ -8,7 +8,6 @@
 
 try:
     import fastapi
-    import rollbar.contrib.fastapi
 
     FASTAPI_INSTALLED = True
 except ImportError:
@@ -28,7 +27,28 @@
 )
 class LoggerMiddlewareTest(BaseTest):
     def setUp(self):
-        importlib.reload(rollbar.contrib.fastapi)
+        importlib.reload(rollbar)
+
+    @mock.patch('rollbar._check_config', return_value=True)
+    @mock.patch('rollbar.send_payload')
+    def test_should_add_framework_version_to_payload(self, mock_send_payload, *mocks):
+        import fastapi
+        from fastapi import FastAPI
+        import rollbar
+        from rollbar.contrib.fastapi.logger import LoggerMiddleware
+
+        self.assertIsNone(rollbar.BASE_DATA_HOOK)
+
+        app = FastAPI()
+        app.add_middleware(LoggerMiddleware)
+
+        rollbar.report_exc_info()
+
+        mock_send_payload.assert_called_once()
+        payload = mock_send_payload.call_args[0][0]
+
+        self.assertIn('fastapi', payload['data']['framework'])
+        self.assertIn(fastapi.__version__, payload['data']['framework'])
 
     def test_should_support_type_hints(self):
         from starlette.types import Receive, Scope, Send

rollbar/test/fastapi_tests/test_middleware.py

@@ -9,7 +9,6 @@
 
 try:
     import fastapi
-    import rollbar.contrib.fastapi
 
     FASTAPI_INSTALLED = True
 except ImportError:
@@ -31,9 +30,9 @@ class ReporterMiddlewareTest(BaseTest):
     default_settings = copy.deepcopy(rollbar.SETTINGS)
 
     def setUp(self):
+        importlib.reload(rollbar)
         rollbar.SETTINGS = copy.deepcopy(self.default_settings)
         rollbar.SETTINGS['handler'] = 'async'
-        importlib.reload(rollbar.contrib.fastapi)
 
     @mock.patch('rollbar.report_exc_info')
     def test_should_catch_and_report_errors(self, mock_report):
@@ -144,6 +143,27 @@ def read_root(path):
             },
         )
 
+    @mock.patch('rollbar._check_config', return_value=True)
+    @mock.patch('rollbar.send_payload')
+    def test_should_add_framework_version_to_payload(self, mock_send_payload, *mocks):
+        import fastapi
+        from fastapi import FastAPI
+        import rollbar
+        from rollbar.contrib.fastapi.middleware import ReporterMiddleware
+
+        self.assertIsNone(rollbar.BASE_DATA_HOOK)
+
+        app = FastAPI()
+        app.add_middleware(ReporterMiddleware)
+
+        rollbar.report_exc_info()
+
+        mock_send_payload.assert_called_once()
+        payload = mock_send_payload.call_args[0][0]
+
+        self.assertIn('fastapi', payload['data']['framework'])
+        self.assertIn(fastapi.__version__, payload['data']['framework'])
+
     @mock.patch('rollbar.lib._async.report_exc_info', new_callable=AsyncMock)
     @mock.patch('rollbar.report_exc_info')
     def test_should_use_async_report_exc_info_if_default_handler(