Refactor integration tests

test/integration/bucket_cleaner.py

@@ -8,11 +8,17 @@
 #
 ######################################################################
 
+import re
+import time
+
 from typing import Optional
 
+from b2sdk.bucket import Bucket
+from b2sdk.file_lock import NO_RETENTION_FILE_SETTING, LegalHold, RetentionMode
+from b2sdk.utils import current_time_millis
 from b2sdk.v2 import *
 
-from .helpers import GENERAL_BUCKET_NAME_PREFIX, BUCKET_CREATED_AT_MILLIS, authorize
+from .helpers import BUCKET_CREATED_AT_MILLIS, GENERAL_BUCKET_NAME_PREFIX, authorize
 
 ONE_HOUR_MILLIS = 60 * 60 * 1000
 
@@ -88,3 +94,61 @@ def cleanup_buckets(self):
                 else:
                     print('Removing bucket:', bucket.name)
                     b2_api.delete_bucket(bucket)
+
+
+def _cleanup_old_buckets(raw_api, api_url, account_auth_token, account_id, bucket_list_dict):
+    for bucket_dict in bucket_list_dict['buckets']:
+        bucket_id = bucket_dict['bucketId']
+        bucket_name = bucket_dict['bucketName']
+        if _should_delete_bucket(bucket_name):
+            print('cleaning up old bucket: ' + bucket_name)
+            _clean_and_delete_bucket(
+                raw_api,
+                api_url,
+                account_auth_token,
+                account_id,
+                bucket_id,
+            )
+
+
+def _clean_and_delete_bucket(raw_api, api_url, account_auth_token, account_id, bucket_id):
+    # Delete the files. This test never creates more than a few files,
+    # so one call to list_file_versions should get them all.
+    versions_dict = raw_api.list_file_versions(api_url, account_auth_token, bucket_id)
+    for version_dict in versions_dict['files']:
+        file_id = version_dict['fileId']
+        file_name = version_dict['fileName']
+        action = version_dict['action']
+        if action in ['hide', 'upload']:
+            print('b2_delete_file', file_name, action)
+            if action == 'upload' and version_dict[
+                'fileRetention'] and version_dict['fileRetention']['value']['mode'] is not None:
+                raw_api.update_file_retention(
+                    api_url,
+                    account_auth_token,
+                    file_id,
+                    file_name,
+                    NO_RETENTION_FILE_SETTING,
+                    bypass_governance=True
+                )
+            raw_api.delete_file_version(api_url, account_auth_token, file_id, file_name)
+        else:
+            print('b2_cancel_large_file', file_name)
+            raw_api.cancel_large_file(api_url, account_auth_token, file_id)
+
+    # Delete the bucket
+    print('b2_delete_bucket', bucket_id)
+    raw_api.delete_bucket(api_url, account_auth_token, account_id, bucket_id)
+
+
+def _should_delete_bucket(bucket_name):
+    # Bucket names for this test look like: c7b22d0b0ad7-1460060364-5670
+    # Other buckets should not be deleted.
+    match = re.match(r'^test-raw-api-[a-f0-9]+-([0-9]+)-([0-9]+)', bucket_name)
+    if match is None:
+        return False
+
+    # Is it more than an hour old?
+    bucket_time = int(match.group(1))
+    now = time.time()
+    return bucket_time + 3600 <= now

test/integration/conftest.py

@@ -8,8 +8,20 @@
 #
 ######################################################################
 
+from io import BytesIO
+from os import environ
+from random import randint
+from time import time
+
 import pytest
 
+from b2sdk.b2http import B2Http
+from b2sdk.encryption.setting import EncryptionAlgorithm, EncryptionMode, EncryptionSetting
+from b2sdk.raw_api import ALL_CAPABILITIES, REALM_URLS, B2RawHTTPApi
+from b2sdk.utils import hex_sha1_of_stream
+
+from .bucket_cleaner import _clean_and_delete_bucket
+
 
 def pytest_addoption(parser):
     """Add a flag for not cleaning up old buckets"""
@@ -20,6 +32,144 @@ def pytest_addoption(parser):
     )
 
 
-@pytest.fixture
+@pytest.fixture(scope='session')
 def dont_cleanup_old_buckets(request):
     return request.config.getoption("--dont-cleanup-old-buckets")
+
+
+@pytest.fixture(scope='session')
+def application_key_id() -> str:
+    key_id = environ.get('B2_TEST_APPLICATION_KEY_ID')
+    assert key_id
+    return key_id
+
+
+@pytest.fixture(scope='session')
+def application_key() -> str:
+    key = environ.get('B2_TEST_APPLICATION_KEY')
+    assert key
+    return key
+
+
+@pytest.fixture(scope='module')
+def raw_api() -> B2RawHTTPApi:
+    return B2RawHTTPApi(B2Http())
+
+
+@pytest.fixture(scope='session')
+def realm_url() -> str:
+    realm = environ.get('B2_TEST_ENVIRONMENT', 'production')
+    return REALM_URLS.get(realm, realm)
+
+
+@pytest.fixture(scope='module')
+def auth_dict(raw_api, application_key, application_key_id, realm_url) -> dict:
+    result = raw_api.authorize_account(realm_url, application_key_id, application_key)
+
+    missing_capabilities = set(ALL_CAPABILITIES) - \
+        {'readBuckets', 'listAllBucketNames'} - set(result['allowed']['capabilities'])
+    assert not missing_capabilities, \
+        f'Seems like a non-full key. Missing capabilities: {missing_capabilities}'
+
+    return result
+
+
+@pytest.fixture(scope='module')
+def account_id(auth_dict) -> str:
+    return auth_dict['accountId']
+
+
+@pytest.fixture(scope='module')
+def account_auth_token(auth_dict) -> str:
+    return auth_dict['authorizationToken']
+
+
+@pytest.fixture(scope='module')
+def api_url(auth_dict) -> str:
+    return auth_dict['apiUrl']
+
+
+@pytest.fixture(scope='module')
+def download_url(auth_dict) -> str:
+    return auth_dict['downloadUrl']
+
+
+@pytest.fixture(scope='module')
+def bucket_dict(raw_api, api_url, account_auth_token, account_id) -> dict:
+    # include the account ID in the bucket name to be
+    # sure it doesn't collide with bucket names from
+    # other accounts.
+    bucket_name = f'test-raw-api-{account_id}-{time():.0f}-{randint(1000, 9999)}'
+
+    bucket = raw_api.create_bucket(
+        api_url,
+        account_auth_token,
+        account_id,
+        bucket_name,
+        'allPublic',
+        is_file_lock_enabled=True,
+    )
+    yield bucket
+
+    _clean_and_delete_bucket(
+        raw_api,
+        api_url,
+        account_auth_token,
+        account_id,
+        bucket['bucketId'],
+    )
+
+
+@pytest.fixture(scope='module')
+def bucket_id(bucket_dict) -> str:
+    return bucket_dict['bucketId']
+
+
+@pytest.fixture(scope='module')
+def bucket_name(bucket_dict) -> str:
+    return bucket_dict['bucketName']
+
+
+@pytest.fixture(scope='module')
+def upload_url_dict(raw_api, api_url, account_auth_token, bucket_id) -> dict:
+    return raw_api.get_upload_url(api_url, account_auth_token, bucket_id)
+
+
+TEST_FILE_NAME = 'test.txt'
+TEST_FILE_CONTENTS = b'hello world'
+
+
+@pytest.fixture(scope='module')
+def file_dict(raw_api, upload_url_dict, sse_b2_aes) -> dict:
+    return raw_api.upload_file(
+        upload_url_dict['uploadUrl'],
+        upload_url_dict['authorizationToken'],
+        TEST_FILE_NAME,
+        len(TEST_FILE_CONTENTS),
+        'text/plain',
+        hex_sha1_of_stream(BytesIO(TEST_FILE_CONTENTS), len(TEST_FILE_CONTENTS)),
+        {'color': 'blue'},
+        BytesIO(TEST_FILE_CONTENTS),
+        server_side_encryption=sse_b2_aes,
+    )
+
+
+@pytest.fixture(scope='module')
+def file_id(file_dict) -> str:
+    return file_dict['fileId']
+
+
+@pytest.fixture(scope='module')
+def download_auth_dict(raw_api, api_url, account_auth_token, bucket_id) -> dict:
+    return raw_api.get_download_authorization(
+        api_url,
+        account_auth_token,
+        bucket_id,
+        TEST_FILE_NAME[:-2],
+        12345,
+    )
+
+
+@pytest.fixture(scope='module')
+def download_auth_token(download_auth_dict) -> str:
+    return download_auth_dict['authorizationToken']