Skip to content
This repository was archived by the owner on Sep 1, 2024. It is now read-only.

Commit abe3fec

Browse files
ramosbugsramosbugs
committed
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.
1 parent 1f6aca3 commit abe3fec

File tree

4 files changed

+242
-46
lines changed

4 files changed

+242
-46
lines changed

‎src/pytest_unflakable/_api.py

Lines changed: 88 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,18 @@
44

55
from __future__ import annotations
66

7+
import gzip
8+
import json
79
import logging
810
import platform
911
import pprint
1012
import sys
1113
import time
12-
from typing import TYPE_CHECKING, Any, List, Mapping, Optional
14+
from typing import TYPE_CHECKING, List, Mapping, Optional
1315

1416
import pkg_resources
1517
import requests
16-
from requests import Response, Session
18+
from requests import HTTPError, Response, Session
1719

1820
if TYPE_CHECKING:
1921
from typing_extensions import NotRequired, TypedDict
@@ -67,39 +69,58 @@ class TestRunRecord(TypedDict):
6769
attempts: List[TestRunAttemptRecord]
6870

6971

70-
class CreateTestSuiteRunRequest(TypedDict):
72+
class CreateTestSuiteRunInlineRequest(TypedDict):
7173
branch: NotRequired[Optional[str]]
7274
commit: NotRequired[Optional[str]]
7375
start_time: str
7476
end_time: str
7577
test_runs: List[TestRunRecord]
7678

7779

80+
class CreateTestSuiteRunUploadRequest(TypedDict):
81+
upload_id: str
82+
83+
84+
class CreateTestSuiteRunUploadUrlResponse(TypedDict):
85+
upload_id: str
86+
87+
7888
class TestSuiteRunPendingSummary(TypedDict):
7989
run_id: str
8090
suite_id: str
8191
branch: NotRequired[Optional[str]]
8292
commit: NotRequired[Optional[str]]
8393

8494

85-
def send_api_request(
86-
api_key: str,
87-
method: Literal['GET', 'POST'],
95+
def __new_requests_session() -> Session:
96+
session = Session()
97+
session.headers['User-Agent'] = USER_AGENT
98+
99+
return session
100+
101+
102+
def __send_api_request(
103+
session: Session,
104+
api_key: Optional[str],
105+
method: Literal['GET', 'POST', 'PUT'],
88106
url: str,
89107
logger: logging.Logger,
90108
headers: Optional[Mapping[str, str | bytes | None]] = None,
91-
json: Optional[Any] = None,
109+
body: Optional[str | bytes] = None,
92110
verify: Optional[bool | str] = None,
93111
) -> Response:
94-
session = Session()
95-
session.headers.update({
96-
'Authorization': f'Bearer {api_key}',
97-
'User-Agent': USER_AGENT,
98-
})
99-
100112
for idx in range(NUM_REQUEST_TRIES):
101113
try:
102-
response = session.request(method, url, headers=headers, json=json, verify=verify)
114+
response = session.request(
115+
method,
116+
url,
117+
headers={
118+
**({'Authorization': f'Bearer {api_key}'} if api_key is not None else {}),
119+
**(headers if headers is not None else {})
120+
},
121+
data=body,
122+
verify=verify,
123+
)
103124
if response.status_code not in [429, 500, 502, 503, 504]:
104125
return response
105126
elif idx + 1 != NUM_REQUEST_TRIES:
@@ -124,7 +145,7 @@ def send_api_request(
124145

125146

126147
def create_test_suite_run(
127-
request: CreateTestSuiteRunRequest,
148+
request: CreateTestSuiteRunInlineRequest,
128149
test_suite_id: str,
129150
api_key: str,
130151
base_url: Optional[str],
@@ -133,7 +154,53 @@ def create_test_suite_run(
133154
) -> TestSuiteRunPendingSummary:
134155
logger.debug(f'creating test suite run {pprint.pformat(request)}')
135156

136-
run_response = send_api_request(
157+
session = __new_requests_session()
158+
159+
create_upload_url_response = __send_api_request(
160+
session=session,
161+
api_key=api_key,
162+
method='POST',
163+
url=(
164+
f'{base_url if base_url is not None else BASE_URL}/api/v1/test-suites/{test_suite_id}'
165+
'/runs/upload'
166+
),
167+
logger=logger,
168+
verify=not insecure_disable_tls_validation,
169+
)
170+
171+
create_upload_url_response.raise_for_status()
172+
if create_upload_url_response.status_code != 201:
173+
raise HTTPError(
174+
f'Expected 201 response but received {create_upload_url_response.status_code}')
175+
176+
upload_presigned_url = create_upload_url_response.headers.get('Location', None)
177+
if upload_presigned_url is None:
178+
raise HTTPError('Location response header not found')
179+
180+
create_upload_url_response_body: CreateTestSuiteRunUploadUrlResponse = (
181+
create_upload_url_response.json()
182+
)
183+
upload_id = create_upload_url_response_body['upload_id']
184+
185+
gzipped_request = gzip.compress(json.dumps(request).encode('utf8'))
186+
upload_response = __send_api_request(
187+
session=session,
188+
api_key=None,
189+
method='PUT',
190+
url=upload_presigned_url,
191+
logger=logger,
192+
headers={
193+
'Content-Encoding': 'gzip',
194+
'Content-Type': 'application/json',
195+
},
196+
body=gzipped_request,
197+
verify=not insecure_disable_tls_validation,
198+
)
199+
upload_response.raise_for_status()
200+
201+
request_body: CreateTestSuiteRunUploadRequest = {'upload_id': upload_id}
202+
run_response = __send_api_request(
203+
session=session,
137204
api_key=api_key,
138205
method='POST',
139206
url=(
@@ -142,7 +209,7 @@ def create_test_suite_run(
142209
),
143210
logger=logger,
144211
headers={'Content-Type': 'application/json'},
145-
json=request,
212+
body=json.dumps(request_body).encode('utf8'),
146213
verify=not insecure_disable_tls_validation,
147214
)
148215
run_response.raise_for_status()
@@ -162,7 +229,10 @@ def get_test_suite_manifest(
162229
) -> TestSuiteManifest:
163230
logger.debug(f'fetching manifest for test suite {test_suite_id}')
164231

165-
manifest_response = send_api_request(
232+
session = __new_requests_session()
233+
234+
manifest_response = __send_api_request(
235+
session=session,
166236
api_key=api_key,
167237
method='GET',
168238
url=(

‎src/pytest_unflakable/_plugin.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import pytest
1515
from _pytest.config import ExitCode
1616

17-
from ._api import (CreateTestSuiteRunRequest, TestAttemptResult,
17+
from ._api import (CreateTestSuiteRunInlineRequest, TestAttemptResult,
1818
TestRunAttemptRecord, TestRunRecord, TestSuiteManifest,
1919
build_test_suite_run_url, create_test_suite_run)
2020

@@ -496,7 +496,7 @@ def pytest_sessionstart(self, session: pytest.Session) -> None:
496496
def _build_test_suite_run_request(
497497
self,
498498
session: pytest.Session,
499-
) -> CreateTestSuiteRunRequest:
499+
) -> CreateTestSuiteRunInlineRequest:
500500
test_runs: List[TestRunRecord] = []
501501
for (test_filename, test_name), item_reports in self.item_reports.items():
502502
is_quarantined = (test_filename, test_name) in self.quarantined_tests
@@ -581,7 +581,7 @@ def _build_test_suite_run_request(
581581
request.update(**({'branch': self.branch} if self.branch is not None else {}))
582582
request.update(**({'commit': self.commit} if self.commit is not None else {}))
583583

584-
return cast(CreateTestSuiteRunRequest, request)
584+
return cast(CreateTestSuiteRunInlineRequest, request)
585585

586586
# Allows us to override the exit code if all the failures are quarantined. We need this to be a
587587
# wrapper so that the default hook still gets invoked and prints the summary line with the test

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy