|
| 1 | +###################################################################### |
| 2 | +# |
| 3 | +# File: test/unit/test_upload_503_token_refresh.py |
| 4 | +# |
| 5 | +# Copyright 2025 Backblaze Inc. All Rights Reserved. |
| 6 | +# |
| 7 | +# License https://www.backblaze.com/using_b2_code.html |
| 8 | +# |
| 9 | +###################################################################### |
| 10 | +""" |
| 11 | +Test to verify that upload tokens are refreshed when a 503 error is returned. |
| 12 | +This test simulates the scenario where a server returns a 503 Service Unavailable |
| 13 | +error during file upload and verifies that the SDK properly requests new upload |
| 14 | +tokens for subsequent retry attempts. |
| 15 | +""" |
| 16 | + |
| 17 | +from __future__ import annotations |
| 18 | + |
| 19 | +from unittest.mock import patch |
| 20 | + |
| 21 | +import pytest |
| 22 | + |
| 23 | +from b2sdk._internal.exception import ServiceError |
| 24 | +from b2sdk.v3 import RawSimulator |
| 25 | +from b2sdk.v3 import B2Api, B2HttpApiConfig, DummyCache, StubAccountInfo |
| 26 | + |
| 27 | +from .test_base import TestBase |
| 28 | + |
| 29 | + |
| 30 | +class TestUpload503TokenRefresh(TestBase): |
| 31 | + """Test that 503 errors trigger upload token refresh using RawSimulator""" |
| 32 | + |
| 33 | + def setUp(self): |
| 34 | + self.bucket_name = 'test-bucket' |
| 35 | + self.account_info = StubAccountInfo() |
| 36 | + self.api = B2Api( |
| 37 | + self.account_info, |
| 38 | + cache=DummyCache(), |
| 39 | + api_config=B2HttpApiConfig(_raw_api_class=RawSimulator), |
| 40 | + ) |
| 41 | + self.simulator = self.api.session.raw_api |
| 42 | + (self.account_id, self.master_key) = self.simulator.create_account() |
| 43 | + self.api.authorize_account( |
| 44 | + application_key_id=self.account_id, |
| 45 | + application_key=self.master_key, |
| 46 | + realm='production', |
| 47 | + ) |
| 48 | + self.bucket = self.api.create_bucket(self.bucket_name, 'allPublic') |
| 49 | + |
| 50 | + def test_upload_503_triggers_token_refresh(self): |
| 51 | + """ |
| 52 | + Test that when a 503 error occurs during upload, the SDK: |
| 53 | + 1. Retries the upload |
| 54 | + 2. Uses a different upload token for the retry |
| 55 | + """ |
| 56 | + # Track upload URLs/tokens used during upload attempts |
| 57 | + upload_urls_used = [] |
| 58 | + |
| 59 | + # Wrap the upload_file method to track upload URLs |
| 60 | + original_upload_file = self.simulator.upload_file |
| 61 | + |
| 62 | + def tracked_upload_file(upload_url, *args, **kwargs): |
| 63 | + upload_urls_used.append(upload_url) |
| 64 | + return original_upload_file(upload_url, *args, **kwargs) |
| 65 | + |
| 66 | + # Inject a 503 error for the first upload attempt |
| 67 | + self.simulator.set_upload_errors( |
| 68 | + [ServiceError('503 service_unavailable Service Unavailable')] |
| 69 | + ) |
| 70 | + |
| 71 | + # Patch upload_file to track URLs |
| 72 | + with patch.object(self.simulator, 'upload_file', side_effect=tracked_upload_file): |
| 73 | + # Perform upload - should fail once with 503, then retry and succeed |
| 74 | + data = b'test data for 503 error scenario' |
| 75 | + file_name = 'test_503.txt' |
| 76 | + |
| 77 | + self.bucket.upload_bytes(data, file_name) |
| 78 | + |
| 79 | + # Verify the upload succeeded |
| 80 | + file_info = self.bucket.get_file_info_by_name(file_name) |
| 81 | + assert file_info is not None |
| 82 | + |
| 83 | + # Verify that at least 2 upload attempts were made |
| 84 | + assert ( |
| 85 | + len(upload_urls_used) >= 2 |
| 86 | + ), f'Expected at least 2 upload attempts, but got {len(upload_urls_used)}' |
| 87 | + |
| 88 | + # Extract auth tokens from the URLs |
| 89 | + # URL format: https://upload.example.com/bucket_id/upload_id/auth_token |
| 90 | + first_url = upload_urls_used[0] |
| 91 | + second_url = upload_urls_used[1] |
| 92 | + |
| 93 | + first_auth_token = first_url.split('/')[-1] |
| 94 | + second_auth_token = second_url.split('/')[-1] |
| 95 | + |
| 96 | + print('\n✓ Upload token refresh test results:') |
| 97 | + print(f' First URL: {first_url}') |
| 98 | + print(f' Second URL: {second_url}') |
| 99 | + print(f' First auth token: {first_auth_token}') |
| 100 | + print(f' Second auth token: {second_auth_token}') |
| 101 | + print(f' Total upload attempts: {len(upload_urls_used)}') |
| 102 | + |
| 103 | + # Verify that auth tokens are different after a 503 error. |
| 104 | + # This confirms that the SDK properly clears cached upload tokens |
| 105 | + # and requests fresh ones when a 503 Service Error occurs. |
| 106 | + assert first_auth_token != second_auth_token, ( |
| 107 | + f'BUG: Auth tokens are the same after 503 error!\n' |
| 108 | + f'Expected different auth tokens, but got:\n' |
| 109 | + f' First: {first_auth_token}\n' |
| 110 | + f' Second: {second_auth_token}\n' |
| 111 | + f'This indicates the SDK is cycling through pre-fetched upload URLs\n' |
| 112 | + f'instead of requesting fresh upload tokens after a 503 error.' |
| 113 | + ) |
0 commit comments