From abe3fecaaee76abc0bb6ecae8ffc80d2470b38df Mon Sep 17 00:00:00 2001 From: "David A. Ramos" Date: Sat, 11 Nov 2023 21:50:18 -0800 Subject: [PATCH] Support test run payloads over 6MB AWS Lambda limits payloads to 6MB, which puts an upper bound on the size of the Unflakable backend API's request bodies. This change leverages a newer backend API endpoint to upload test runs to presigned S3 URLs, which are not subject to AWS Lambda limits. The backend then retrieves the upload directly from S3. --- src/pytest_unflakable/_api.py | 106 +++++++++++++++---- src/pytest_unflakable/_plugin.py | 6 +- tests/common.py | 173 ++++++++++++++++++++++++++----- tests/test_unflakable.py | 3 + 4 files changed, 242 insertions(+), 46 deletions(-) diff --git a/src/pytest_unflakable/_api.py b/src/pytest_unflakable/_api.py index 635ee2c..bce77a5 100644 --- a/src/pytest_unflakable/_api.py +++ b/src/pytest_unflakable/_api.py @@ -4,16 +4,18 @@ from __future__ import annotations +import gzip +import json import logging import platform import pprint import sys import time -from typing import TYPE_CHECKING, Any, List, Mapping, Optional +from typing import TYPE_CHECKING, List, Mapping, Optional import pkg_resources import requests -from requests import Response, Session +from requests import HTTPError, Response, Session if TYPE_CHECKING: from typing_extensions import NotRequired, TypedDict @@ -67,7 +69,7 @@ class TestRunRecord(TypedDict): attempts: List[TestRunAttemptRecord] -class CreateTestSuiteRunRequest(TypedDict): +class CreateTestSuiteRunInlineRequest(TypedDict): branch: NotRequired[Optional[str]] commit: NotRequired[Optional[str]] start_time: str @@ -75,6 +77,14 @@ class CreateTestSuiteRunRequest(TypedDict): test_runs: List[TestRunRecord] +class CreateTestSuiteRunUploadRequest(TypedDict): + upload_id: str + + +class CreateTestSuiteRunUploadUrlResponse(TypedDict): + upload_id: str + + class TestSuiteRunPendingSummary(TypedDict): run_id: str suite_id: str @@ -82,24 +92,35 @@ class TestSuiteRunPendingSummary(TypedDict): commit: NotRequired[Optional[str]] -def send_api_request( - api_key: str, - method: Literal['GET', 'POST'], +def __new_requests_session() -> Session: + session = Session() + session.headers['User-Agent'] = USER_AGENT + + return session + + +def __send_api_request( + session: Session, + api_key: Optional[str], + method: Literal['GET', 'POST', 'PUT'], url: str, logger: logging.Logger, headers: Optional[Mapping[str, str | bytes | None]] = None, - json: Optional[Any] = None, + body: Optional[str | bytes] = None, verify: Optional[bool | str] = None, ) -> Response: - session = Session() - session.headers.update({ - 'Authorization': f'Bearer {api_key}', - 'User-Agent': USER_AGENT, - }) - for idx in range(NUM_REQUEST_TRIES): try: - response = session.request(method, url, headers=headers, json=json, verify=verify) + response = session.request( + method, + url, + headers={ + **({'Authorization': f'Bearer {api_key}'} if api_key is not None else {}), + **(headers if headers is not None else {}) + }, + data=body, + verify=verify, + ) if response.status_code not in [429, 500, 502, 503, 504]: return response elif idx + 1 != NUM_REQUEST_TRIES: @@ -124,7 +145,7 @@ def send_api_request( def create_test_suite_run( - request: CreateTestSuiteRunRequest, + request: CreateTestSuiteRunInlineRequest, test_suite_id: str, api_key: str, base_url: Optional[str], @@ -133,7 +154,53 @@ def create_test_suite_run( ) -> TestSuiteRunPendingSummary: logger.debug(f'creating test suite run {pprint.pformat(request)}') - run_response = send_api_request( + session = __new_requests_session() + + create_upload_url_response = __send_api_request( + session=session, + api_key=api_key, + method='POST', + url=( + f'{base_url if base_url is not None else BASE_URL}/api/v1/test-suites/{test_suite_id}' + '/runs/upload' + ), + logger=logger, + verify=not insecure_disable_tls_validation, + ) + + create_upload_url_response.raise_for_status() + if create_upload_url_response.status_code != 201: + raise HTTPError( + f'Expected 201 response but received {create_upload_url_response.status_code}') + + upload_presigned_url = create_upload_url_response.headers.get('Location', None) + if upload_presigned_url is None: + raise HTTPError('Location response header not found') + + create_upload_url_response_body: CreateTestSuiteRunUploadUrlResponse = ( + create_upload_url_response.json() + ) + upload_id = create_upload_url_response_body['upload_id'] + + gzipped_request = gzip.compress(json.dumps(request).encode('utf8')) + upload_response = __send_api_request( + session=session, + api_key=None, + method='PUT', + url=upload_presigned_url, + logger=logger, + headers={ + 'Content-Encoding': 'gzip', + 'Content-Type': 'application/json', + }, + body=gzipped_request, + verify=not insecure_disable_tls_validation, + ) + upload_response.raise_for_status() + + request_body: CreateTestSuiteRunUploadRequest = {'upload_id': upload_id} + run_response = __send_api_request( + session=session, api_key=api_key, method='POST', url=( @@ -142,7 +209,7 @@ def create_test_suite_run( ), logger=logger, headers={'Content-Type': 'application/json'}, - json=request, + body=json.dumps(request_body).encode('utf8'), verify=not insecure_disable_tls_validation, ) run_response.raise_for_status() @@ -162,7 +229,10 @@ def get_test_suite_manifest( ) -> TestSuiteManifest: logger.debug(f'fetching manifest for test suite {test_suite_id}') - manifest_response = send_api_request( + session = __new_requests_session() + + manifest_response = __send_api_request( + session=session, api_key=api_key, method='GET', url=( diff --git a/src/pytest_unflakable/_plugin.py b/src/pytest_unflakable/_plugin.py index b6ab3ea..6bea264 100644 --- a/src/pytest_unflakable/_plugin.py +++ b/src/pytest_unflakable/_plugin.py @@ -14,7 +14,7 @@ import pytest from _pytest.config import ExitCode -from ._api import (CreateTestSuiteRunRequest, TestAttemptResult, +from ._api import (CreateTestSuiteRunInlineRequest, TestAttemptResult, TestRunAttemptRecord, TestRunRecord, TestSuiteManifest, build_test_suite_run_url, create_test_suite_run) @@ -496,7 +496,7 @@ def pytest_sessionstart(self, session: pytest.Session) -> None: def _build_test_suite_run_request( self, session: pytest.Session, - ) -> CreateTestSuiteRunRequest: + ) -> CreateTestSuiteRunInlineRequest: test_runs: List[TestRunRecord] = [] for (test_filename, test_name), item_reports in self.item_reports.items(): is_quarantined = (test_filename, test_name) in self.quarantined_tests @@ -581,7 +581,7 @@ def _build_test_suite_run_request( request.update(**({'branch': self.branch} if self.branch is not None else {})) request.update(**({'commit': self.commit} if self.commit is not None else {})) - return cast(CreateTestSuiteRunRequest, request) + return cast(CreateTestSuiteRunInlineRequest, request) # Allows us to override the exit code if all the failures are quarantined. We need this to be a # wrapper so that the default hook still gets invoked and prints the summary line with the test diff --git a/tests/common.py b/tests/common.py index d54dd8f..88dc8e2 100644 --- a/tests/common.py +++ b/tests/common.py @@ -1,7 +1,10 @@ """Tests for pytest_unflakable plugin.""" - +import gzip +import hashlib # Copyright (c) 2022-2023 Developer Innovations, LLC import itertools +import json +import os import re import subprocess from enum import Enum @@ -27,6 +30,7 @@ MOCK_RUN_ID = 'MOCK_RUN_ID' MOCK_SUITE_ID = 'MOCK_SUITE_ID' +MOCK_TEAM_ID = 'MOCK_TEAM_ID' # e.g., 2022-01-23T04:05:06.000000+00:00 TIMESTAMP_REGEX = ( @@ -229,16 +233,72 @@ def mock_run( raise RuntimeError(f'unexpected git call with args: {repr(args)}') -def mock_create_test_suite_run_response( - request: requests.Request, +__uploads: Dict[str, Optional[_api.CreateTestSuiteRunInlineRequest]] = {} + + +def __upload_id_for_current_test() -> str: + return hashlib.sha1(os.environ['PYTEST_CURRENT_TEST'].encode('utf8')).hexdigest() + + +def __upload_url(upload_id: str) -> str: + return ( + f'https://s3.mock.amazonaws.com/unflakable-backend-mock-test-uploads/teams/{MOCK_TEAM_ID}' + f'/suites/{MOCK_SUITE_ID}/runs/upload/{upload_id}?X-Amz-Signature=MOCK_SIGNATURE' + ) + + +def __mock_create_test_suite_run_upload_url_response( + upload_id: str, + request: requests_mock.request._RequestObjectProxy, + context: requests_mock.response._Context, +) -> _api.CreateTestSuiteRunUploadUrlResponse: + upload_url = __upload_url(upload_id) + assert upload_url not in __uploads + __uploads[upload_url] = None + + context.headers['Location'] = upload_url + return { + 'upload_id': upload_id, + } + + +def __match_upload(request: requests_mock.request._RequestObjectProxy) -> bool: + return re.match( + r'^%s[0-9a-f]{40}%s$' % ( + re.escape( + 'https://s3.mock.amazonaws.com/unflakable-backend-mock-test-uploads/teams/' + f'{MOCK_TEAM_ID}/suites/{MOCK_SUITE_ID}/runs/upload/' + ), + re.escape('?X-Amz-Signature=MOCK_SIGNATURE') + ), + request.url + ) is not None + + +def __mock_upload_response( + request: requests_mock.request._RequestObjectProxy, + context: requests_mock.response._Context, +) -> bytes: + assert request.url in __uploads + assert __uploads[request.url] is None, 'duplicate upload' + __uploads[request.url] = json.loads(gzip.decompress(request.body)) + return b'' + + +def __mock_create_test_suite_run_response( + request: requests_mock.request._RequestObjectProxy, context: requests_mock.response._Context, ) -> _api.TestSuiteRunPendingSummary: - request_body: _api.CreateTestSuiteRunRequest = request.json() + request_body: _api.CreateTestSuiteRunUploadRequest = request.json() + upload_url = __upload_url(request_body['upload_id']) + upload = __uploads[upload_url] + assert upload is not None, 'missing upload' + return { 'run_id': MOCK_RUN_ID, 'suite_id': MOCK_SUITE_ID, - 'branch': request_body.get('branch'), - 'commit': request_body.get('commit'), + 'branch': upload.get('branch'), + 'commit': upload.get('commit'), } @@ -274,7 +334,7 @@ def run_test_case( ) -> None: api_key_path = pytester.makefile('', expected_api_key) if use_api_key_path else None requests_mocker.get( - url='https://app.unflakable.com/api/v1/test-suites/MOCK_SUITE_ID/manifest', + url=f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/manifest', request_headers={'Authorization': f'Bearer {expected_api_key}'}, complete_qs=True, response_list=[ @@ -286,11 +346,11 @@ def run_test_case( }] if manifest is not None else []) ) + upload_id = __upload_id_for_current_test() requests_mocker.post( - url='https://app.unflakable.com/api/v1/test-suites/MOCK_SUITE_ID/runs', + url=f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/runs/upload', request_headers={ - 'Authorization': f'Bearer {expected_api_key}', - 'Content-Type': 'application/json', + 'Authorization': f'Bearer {expected_api_key}' }, complete_qs=True, response_list=[ @@ -298,10 +358,38 @@ def run_test_case( for _ in range(failed_upload_requests) ] + [{ 'status_code': 201, - 'json': mock_create_test_suite_run_response, + 'json': lambda request, context: __mock_create_test_suite_run_upload_url_response( + upload_id, + request, + context, + ), }] ) + requests_mocker.put( + # The __match_upload() function matches the URL. + requests_mock.ANY, + request_headers={ + 'Content-Encoding': 'gzip', + 'Content-Type': 'application/json', + }, + complete_qs=True, + status_code=200, + additional_matcher=__match_upload, + content=__mock_upload_response, + ) + + requests_mocker.post( + url=f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/runs', + request_headers={ + 'Authorization': f'Bearer {expected_api_key}', + 'Content-Type': 'application/json', + }, + complete_qs=True, + status_code=201, + json=__mock_create_test_suite_run_response, + ) + pytest_args: List[str] = ( (['--enable-unflakable'] if plugin_enabled else []) + (['--api-key-path', str(api_key_path)] if api_key_path is not None else []) + @@ -314,6 +402,7 @@ def run_test_case( ) + list(extra_args) ) + __pytest_current_test = os.environ['PYTEST_CURRENT_TEST'] if monkeypatch is not None: with monkeypatch.context() as mp: for key, val in (env_vars if env_vars is not None else {}).items(): @@ -323,6 +412,9 @@ def run_test_case( else: result = pytester.runpytest(*pytest_args) + # pytester clears PYTEST_CURRENT_TEST for some reason. + os.environ['PYTEST_CURRENT_TEST'] = __pytest_current_test + if verbose: test_outcomes_output = [ # Per-file test outcomes (one line for each test, color-coded). @@ -508,9 +600,10 @@ def run_test_case( request = requests_mocker.request_history[manifest_attempt] assert request.url == ( - 'https://app.unflakable.com/api/v1/test-suites/MOCK_SUITE_ID/manifest' + f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/manifest' ) assert request.method == 'GET' + assert request.headers.get('Authorization', '') == f'Bearer {expected_api_key}' assert request.body is None if manifest_attempt > 0: @@ -527,26 +620,57 @@ def run_test_case( ) for upload_attempt in range(expected_upload_attempts): - create_test_suite_run_request = requests_mocker.request_history[ + create_upload_url_request = requests_mocker.request_history[ expected_get_test_suite_manifest_attempts + upload_attempt ] - assert create_test_suite_run_request.url == ( - 'https://app.unflakable.com/api/v1/test-suites/MOCK_SUITE_ID/runs') - assert create_test_suite_run_request.method == 'POST' + assert create_upload_url_request.url == ( + f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/runs/upload') + assert create_upload_url_request.method == 'POST' + assert ( + create_upload_url_request.headers.get('Authorization') + == f'Bearer {expected_api_key}' + ) - create_test_suite_run_body: _api.CreateTestSuiteRunRequest = ( - create_test_suite_run_request.json() + # Failed attempts only include the initial request. + if upload_attempt < failed_upload_requests: + continue + + upload_request = requests_mocker.request_history[ + expected_get_test_suite_manifest_attempts + upload_attempt + 1 + ] + assert upload_request.url == __upload_url(upload_id) + assert upload_request.method == 'PUT' + assert upload_request.headers.get('Content-Encoding') == 'gzip' + assert upload_request.headers.get('Content-Type') == 'application/json' + upload_body: _api.CreateTestSuiteRunInlineRequest = ( + json.loads(gzip.decompress(upload_request.body)) ) if expected_commit is not None: - assert create_test_suite_run_body['commit'] == expected_commit + assert upload_body['commit'] == expected_commit else: - assert 'commit' not in create_test_suite_run_body + assert 'commit' not in upload_body if expected_branch is not None: - assert create_test_suite_run_body['branch'] == expected_branch + assert upload_body['branch'] == expected_branch else: - assert 'branch' not in create_test_suite_run_body + assert 'branch' not in upload_body + + create_test_suite_run_request = requests_mocker.request_history[ + expected_get_test_suite_manifest_attempts + upload_attempt + 2 + ] + assert create_test_suite_run_request.url == ( + f'https://app.unflakable.com/api/v1/test-suites/{MOCK_SUITE_ID}/runs') + assert create_test_suite_run_request.method == 'POST' + assert ( + create_test_suite_run_request.headers.get('Authorization') + == f'Bearer {expected_api_key}' + ) + assert create_test_suite_run_request.headers.get('Content-Type') == 'application/json' + create_test_suite_run_body: _api.CreateTestSuiteRunUploadRequest = ( + create_test_suite_run_request.json() + ) + assert create_test_suite_run_body['upload_id'] == upload_id if upload_attempt > 0: assert ( @@ -557,7 +681,8 @@ def run_test_case( ) assert requests_mocker.call_count == ( - expected_get_test_suite_manifest_attempts + expected_upload_attempts + expected_get_test_suite_manifest_attempts + + failed_upload_requests + 3 * (expected_upload_attempts - failed_upload_requests) ), 'Expected %d total API requests, but received %d' % ( expected_get_test_suite_manifest_attempts + expected_upload_attempts, requests_mocker.call_count, @@ -566,8 +691,6 @@ def run_test_case( # Checked expected User-Agent. We do this here instead of using an `additional_matcher` to # make errors easier to diagnose. for request in requests_mocker.request_history: - assert request.headers.get('Authorization', '') == f'Bearer {expected_api_key}' - assert_regex( r'^unflakable-pytest-plugin/.* \(PyTest .*; Python .*; Platform .*\)$', request.headers.get('User-Agent', '') diff --git a/tests/test_unflakable.py b/tests/test_unflakable.py index 6def7c7..f39f4bd 100644 --- a/tests/test_unflakable.py +++ b/tests/test_unflakable.py @@ -1949,6 +1949,9 @@ def test_pass2(): expect_progress=False, ) + # Prevent duplicate upload error. + os.environ['PYTEST_CURRENT_TEST'] += '-step2' + run_test_case( pytester, manifest,