Update test framework for starlette 0.15.0 (#1174)

* Update test framework for starlette 0.15.0

* Migrate to ASGI middleware for starlette

* CHANGELOG

* Don't capture response body

Turns out we don't do this in other frameworks.

Author: basepi

.ci/.jenkins_exclude.yml

@@ -204,6 +204,10 @@ exclude:
     FRAMEWORK: starlette-0.13
   - PYTHON_VERSION: python-3.6
     FRAMEWORK: starlette-0.13
+  - PYTHON_VERSION: pypy-3
+    FRAMEWORK: starlette-0.14
+  - PYTHON_VERSION: python-3.6
+    FRAMEWORK: starlette-0.14
   - PYTHON_VERSION: pypy-3
     FRAMEWORK: starlette-newest
   - PYTHON_VERSION: python-3.6

.ci/.jenkins_framework_full.yml

@@ -70,7 +70,8 @@ FRAMEWORK:
   - asyncpg-newest
   - tornado-newest
   - starlette-0.13
-  # - starlette-newest # disabled for now, see https://github.com/elastic/apm-agent-python/issues/1172
+  - starlette-0.14
+  - starlette-newest
   - pymemcache-3.0
   - pymemcache-newest
   - graphene-2

CHANGELOG.asciidoc

@@ -42,6 +42,7 @@ endif::[]
 [float]
 ===== Bug fixes
 
+* Fix for Starlette 0.15.0 error collection {pull}1174[#1174]
 * Fix for Starlette static files {pull}1137[#1137]
 
 [[release-notes-6.x]]

elasticapm/contrib/starlette/__init__.py

@@ -31,14 +31,14 @@
 
 from __future__ import absolute_import
 
+import asyncio
+import functools
 from typing import Dict, Optional
 
 import starlette
-from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
-from starlette.responses import Response
 from starlette.routing import Match, Mount
-from starlette.types import ASGIApp
+from starlette.types import ASGIApp, Message
 
 import elasticapm
 import elasticapm.instrumentation.control
@@ -70,32 +70,28 @@ def make_apm_client(config: Optional[Dict] = None, client_cls=Client, **defaults
     return client_cls(config, **defaults)
 
 
-class ElasticAPM(BaseHTTPMiddleware):
+class ElasticAPM:
     """
     Starlette / FastAPI middleware for Elastic APM capturing.
 
-    >>> elasticapm = make_apm_client({
+    >>> apm = make_apm_client({
         >>> 'SERVICE_NAME': 'myapp',
         >>> 'DEBUG': True,
         >>> 'SERVER_URL': 'http://localhost:8200',
         >>> 'CAPTURE_HEADERS': True,
         >>> 'CAPTURE_BODY': 'all'
     >>> })
 
-    >>> app.add_middleware(ElasticAPM, client=elasticapm)
+    >>> app.add_middleware(ElasticAPM, client=apm)
 
     Pass an arbitrary APP_NAME and SECRET_TOKEN::
 
     >>> elasticapm = ElasticAPM(app, service_name='myapp', secret_token='asdasdasd')
 
-    Pass an explicit client::
+    Pass an explicit client (don't pass in additional options in this case)::
 
     >>> elasticapm = ElasticAPM(app, client=client)
 
-    Automatically configure logging::
-
-    >>> elasticapm = ElasticAPM(app, logging=True)
-
     Capture an exception::
 
     >>> try:
@@ -108,34 +104,69 @@ class ElasticAPM(BaseHTTPMiddleware):
     >>> elasticapm.capture_message('hello, world!')
     """
 
-    def __init__(self, app: ASGIApp, client: Client):
+    def __init__(self, app: ASGIApp, client: Optional[Client], **kwargs):
         """
 
         Args:
             app (ASGIApp): Starlette app
             client (Client): ElasticAPM Client
         """
-        self.client = client
+        if client:
+            self.client = client
+        else:
+            self.client = make_apm_client(**kwargs)
 
         if self.client.config.instrument and self.client.config.enabled:
             elasticapm.instrumentation.control.instrument()
 
-        super().__init__(app)
-
-    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
-        """Processes the whole request APM capturing.
+        # If we ever make this a general-use ASGI middleware we should use
+        # `asgiref.conpatibility.guarantee_single_callable(app)` here
+        self.app = app
 
+    async def __call__(self, scope, receive, send):
+        """
         Args:
-            request (Request)
-            call_next (RequestResponseEndpoint): Next request process in Starlette.
-
-        Returns:
-            Response
+            scope: ASGI scope dictionary
+            receive: receive awaitable callable
+            send: send awaitable callable
         """
+
+        @functools.wraps(send)
+        async def wrapped_send(message):
+            if message.get("type") == "http.response.start":
+                await set_context(
+                    lambda: get_data_from_response(message, self.client.config, constants.TRANSACTION), "response"
+                )
+                result = "HTTP {}xx".format(message["status"] // 100)
+                elasticapm.set_transaction_result(result, override=False)
+            await send(message)
+
+        # When we consume the body from receive, we replace the streaming
+        # mechanism with a mocked version -- this workaround came from
+        # https://github.com/encode/starlette/issues/495#issuecomment-513138055
+        body = b""
+        while True:
+            message = await receive()
+            if not message:
+                break
+            if message["type"] == "http.request":
+                b = message.get("body", b"")
+                if b:
+                    body += b
+                if not message.get("more_body", False):
+                    break
+            if message["type"] == "http.disconnect":
+                break
+
+        async def _receive() -> Message:
+            await asyncio.sleep(0)
+            return {"type": "http.request", "body": body}
+
+        request = Request(scope, receive=_receive)
         await self._request_started(request)
 
         try:
-            response = await call_next(request)
+            await self.app(scope, _receive, wrapped_send)
             elasticapm.set_transaction_outcome(constants.OUTCOME.SUCCESS, override=False)
         except Exception:
             await self.capture_exception(
@@ -146,13 +177,9 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -
             elasticapm.set_context({"status_code": 500}, "response")
 
             raise
-        else:
-            await self._request_finished(response)
         finally:
             self.client.end_transaction()
 
-        return response
-
     async def capture_exception(self, *args, **kwargs):
         """Captures your exception.
 
@@ -195,19 +222,6 @@ async def _request_started(self, request: Request):
             transaction_name = self.get_route_name(request) or request.url.path
             elasticapm.set_transaction_name("{} {}".format(request.method, transaction_name), override=False)
 
-    async def _request_finished(self, response: Response):
-        """Captures the end of the request processing to APM.
-
-        Args:
-            response (Response)
-        """
-        await set_context(
-            lambda: get_data_from_response(response, self.client.config, constants.TRANSACTION), "response"
-        )
-
-        result = "HTTP {}xx".format(response.status_code // 100)
-        elasticapm.set_transaction_result(result, override=False)
-
     def get_route_name(self, request: Request) -> str:
         app = request.app
         scope = request.scope

elasticapm/contrib/starlette/utils.py

@@ -29,8 +29,8 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 import asyncio
 
+from starlette.datastructures import Headers
 from starlette.requests import Request
-from starlette.responses import Response
 from starlette.types import Message
 
 from elasticapm.conf import Config, constants
@@ -73,11 +73,11 @@ async def get_data_from_request(request: Request, config: Config, event_type: st
     return result
 
 
-async def get_data_from_response(response: Response, config: Config, event_type: str) -> dict:
+async def get_data_from_response(message: dict, config: Config, event_type: str) -> dict:
     """Loads data from response for APM capturing.
 
     Args:
-        response (Response)
+        message (dict)
         config (Config)
         event_type (str)
 
@@ -86,16 +86,13 @@ async def get_data_from_response(response: Response, config: Config, event_type:
     """
     result = {}
 
-    if isinstance(getattr(response, "status_code", None), compat.integer_types):
-        result["status_code"] = response.status_code
+    if "status_code" in message:
+        result["status_code"] = message["status"]
 
-    if config.capture_headers and getattr(response, "headers", None):
-        headers = response.headers
+    if config.capture_headers and "headers" in message:
+        headers = Headers(raw=message["headers"])
         result["headers"] = {key: ";".join(headers.getlist(key)) for key in compat.iterkeys(headers)}
 
-    if config.capture_body in ("all", event_type) and hasattr(response, "body"):
-        result["body"] = response.body.decode("utf-8")
-
     return result
 
 

tests/contrib/asyncio/starlette_tests.py

@@ -66,6 +66,7 @@ def app(elasticapm_client):
 
     @app.route("/", methods=["GET", "POST"])
     async def hi(request):
+        await request.body()
         with async_capture_span("test"):
             pass
         return PlainTextResponse("ok")

tests/requirements/reqs-starlette-0.13.txt

@@ -1,4 +1,4 @@
-starlette==0.13.0
+starlette>=0.13,<0.14
 aiofiles==0.7.0
 requests==2.23.0
 -r reqs-base.txt

tests/requirements/reqs-starlette-0.14.txt

@@ -0,0 +1,3 @@
+starlette>=0.14,<0.15
+requests==2.23.0
+-r reqs-base.txt

tests/requirements/reqs-starlette-newest.txt

@@ -1,4 +1,4 @@
-starlette
+starlette>=0.15
 aiofiles
 requests
 flask