Retry integration tests on transient service errors

b2sdk/_internal/exception.py

@@ -433,6 +433,15 @@ class ServiceError(TransientErrorMixin, B2Error):
     Used for HTTP status codes 500 through 599.
     """
 
+    def __init__(self, status, code, message):
+        super().__init__()
+        self._status = status
+        self._code = code
+        self._message = message
+
+    def __str__(self):
+        return f'{self._status} {self._code} {self._message}'
+
 
 class CapExceeded(B2Error):
     def __str__(self):
@@ -744,5 +753,5 @@ def interpret_b2_error(
     elif status == 429:
         return TooManyRequests(retry_after_seconds=response_headers.get('retry-after'))
     elif 500 <= status < 600:
-        return ServiceError('%d %s %s' % (status, code, message))
+        return ServiceError(status, code, message)
     return UnknownError('%d %s %s' % (status, code, message))

b2sdk/_internal/testing/helpers/bucket_manager.py

@@ -25,6 +25,7 @@
     BucketIdNotFound,
     DuplicateBucketName,
     FileNotPresent,
+    ServiceError,
     TooManyRequests,
 )
 from b2sdk._internal.file_lock import NO_RETENTION_FILE_SETTING, LegalHold, RetentionMode
@@ -43,6 +44,12 @@
 logger = logging.getLogger(__name__)
 
 
+def _retry_bucket_test_operation(exc: BaseException) -> bool:
+    return isinstance(exc, TooManyRequests) or (
+        isinstance(exc, ServiceError) and exc._status == 503
+    )
+
+
 class BucketManager:
     def __init__(
         self,
@@ -83,7 +90,7 @@ def new_bucket_info(self) -> dict:
         }
 
     @tenacity.retry(
-        retry=tenacity.retry_if_exception_type(TooManyRequests),
+        retry=tenacity.retry_if_exception(_retry_bucket_test_operation),
         wait=tenacity.wait_exponential(),
         stop=tenacity.stop_after_attempt(8),
     )
@@ -152,7 +159,7 @@ def clean_buckets(self, quick=False):
             print(bucket)
 
     @tenacity.retry(
-        retry=tenacity.retry_if_exception_type(TooManyRequests),
+        retry=tenacity.retry_if_exception(_retry_bucket_test_operation),
         wait=tenacity.wait_exponential(),
         stop=tenacity.stop_after_attempt(8),
     )

changelog.d/1232.changed.md

@@ -0,0 +1 @@
+Retry integration tests automatically when they fail because of a transient HTTP 503 from B2.

test/integration/conftest.py

@@ -7,9 +7,44 @@
 # License https://www.backblaze.com/using_b2_code.html
 #
 ######################################################################
+from __future__ import annotations
+
 import pytest
 
+from b2sdk._internal.exception import ServiceError, TooManyRequests
+
+RETRYABLE_SERVICE_ERROR_STATUSES = {500, 503}
+INTEGRATION_TEST_RETRY_COUNT = 4
+
 
 @pytest.fixture(scope='session', autouse=True)
 def auto_change_account_info_dir(change_account_info_dir):
     pass
+
+
+@pytest.hookimpl(tryfirst=True)
+def pytest_pyfunc_call(pyfuncitem):
+    testfunction = pyfuncitem.obj
+    funcargs = pyfuncitem.funcargs
+    testargs = {arg: funcargs[arg] for arg in pyfuncitem._fixtureinfo.argnames}
+
+    for attempt in range(INTEGRATION_TEST_RETRY_COUNT + 1):
+        try:
+            testfunction(**testargs)
+            return True
+        except ServiceError as exc:
+            if exc._status not in RETRYABLE_SERVICE_ERROR_STATUSES:
+                raise
+            if attempt >= INTEGRATION_TEST_RETRY_COUNT:
+                raise
+            print(
+                f'Retrying {pyfuncitem.nodeid} after transient service error {exc._status}:'
+                f' attempt {attempt + 1} of {INTEGRATION_TEST_RETRY_COUNT}'
+            )
+        except TooManyRequests:
+            if attempt >= INTEGRATION_TEST_RETRY_COUNT:
+                raise
+            print(
+                f'Retrying {pyfuncitem.nodeid} after transient too many requests:'
+                f' attempt {attempt + 1} of {INTEGRATION_TEST_RETRY_COUNT}'
+            )

test/integration/test_upload.py

@@ -10,15 +10,27 @@
 from __future__ import annotations
 
 import io
+import logging
+import secrets
 
-from b2sdk._internal.b2http import B2Http
+from b2sdk._internal.b2http import B2Http, HttpCallback
 from b2sdk._internal.encryption.setting import EncryptionKey, EncryptionSetting
 from b2sdk._internal.encryption.types import EncryptionAlgorithm, EncryptionMode
+from b2sdk._internal.utils import hex_sha1_of_stream
 from b2sdk.v2 import B2RawHTTPApi
 from b2sdk.v3.testing import IntegrationTestBase
 
 from .test_raw_api import authorize_raw_api
 
+logger = logging.getLogger(__name__)
+
+
+class FailSomeUploads(HttpCallback):
+    def pre_request(self, method, url, headers):
+        if method == 'POST' and 'b2_upload_file' in url:
+            headers['X-Bz-Test-Mode'] = 'fail_some_uploads'
+            logger.info('Added X-Bz-Test-Mode=fail_some_uploads header to %s', url)
+
 
 class TestUnboundStreamUpload(IntegrationTestBase):
     def assert_data_uploaded_via_stream(self, data: bytes, part_size: int | None = None):
@@ -46,6 +58,34 @@ def test_streamed_large_buffer_small_part_size(self):
 
 
 class TestUploadLargeFile(IntegrationTestBase):
+    def test_raw_upload_with_intermittent_failures(self):
+        bucket = self.create_bucket()
+        raw_api = self.b2_api.session.raw_api
+        b2_http = raw_api.b2_http
+        account_info = self.b2_api.account_info
+        callback = FailSomeUploads()
+        b2_http.add_callback(callback)
+        try:
+            payload = b'payload'
+            file_name = f'fail-some-uploads-{secrets.token_hex(4)}'
+            upload_url = raw_api.get_upload_url(
+                account_info.get_api_url(),
+                account_info.get_account_auth_token(),
+                bucket.id_,
+            )
+            raw_api.upload_file(
+                upload_url['uploadUrl'],
+                upload_url['authorizationToken'],
+                file_name,
+                len(payload),
+                'text/plain',
+                hex_sha1_of_stream(io.BytesIO(payload), len(payload)),
+                {},
+                io.BytesIO(payload),
+            )
+        finally:
+            b2_http.callbacks.remove(callback)
+
     def test_ssec_key_id(self):
         sse_c = EncryptionSetting(
             mode=EncryptionMode.SSE_C,

test/unit/test_exception.py

@@ -159,6 +159,9 @@ def test_bad_bucket_id(self):
     def test_service_error(self):
         error = interpret_b2_error(500, 'code', 'message', {})
         assert isinstance(error, ServiceError)
+        assert error._status == 500
+        assert error._code == 'code'
+        assert error._message == 'message'
         assert '500 code message' == str(error)
 
     def test_unknown_error(self):