Add common logic for creating and cleaning buckets in tests

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Infrastructure
+* Wrapper for B2Api class which can be used for test purposes
+
 ## [1.21.0] - 2023-04-17
 
 ### Added

b2sdk/test/__init__.py

@@ -0,0 +1,9 @@
+######################################################################
+#
+# File: b2sdk/test/__init__.py
+#
+# Copyright 2023 Backblaze Inc. All Rights Reserved.
+#
+# License https://www.backblaze.com/using_b2_code.html
+#
+######################################################################

b2sdk/test/api_test_manager.py

@@ -0,0 +1,166 @@
+######################################################################
+#
+# File: b2sdk/test/api_test_manager.py
+#
+# Copyright 2023 Backblaze Inc. All Rights Reserved.
+#
+# License https://www.backblaze.com/using_b2_code.html
+#
+######################################################################
+import time
+import uuid
+
+from datetime import datetime
+from os import environ
+from typing import Union
+
+import backoff
+
+from .bucket_tracking import BucketTrackingMixin
+from b2sdk.v2 import (
+    NO_RETENTION_FILE_SETTING, B2Api, Bucket, InMemoryAccountInfo, InMemoryCache, LegalHold,
+    RetentionMode
+)
+from b2sdk.v2.exception import BucketIdNotFound, DuplicateBucketName, FileNotPresent, TooManyRequests, NonExistentBucket
+
+SHORT_SHA = environ.get('GITHUB_SHA', 'local')[:10]
+BUCKET_NAME_PREFIX = f"b2test-{SHORT_SHA}"
+
+
+def generate_bucket_name() -> str:
+    return f"{BUCKET_NAME_PREFIX}-{uuid.uuid4()}"
+
+
+def current_time_millis() -> int:
+    return int(round(time.time() * 1000))
+
+
+class ApiTestManager(BucketTrackingMixin, B2Api):
+    """
+    B2Api wrapper which should only be used for testing purposes!
+    """
+
+    def __init__(self, account_id: str, application_key: str, realm: str, *args, **kwargs):
+        info = InMemoryAccountInfo()
+        cache = InMemoryCache()
+        super().__init__(info, cache=cache, *args, **kwargs)
+        self.authorize_account(realm, account_id, application_key)
+
+    @backoff.on_exception(
+        backoff.constant,
+        DuplicateBucketName,
+        max_tries=8,
+    )
+    def create_test_bucket(self, bucket_type="allPublic", **kwargs) -> Bucket:
+        bucket_name = generate_bucket_name()
+        print(f'Creating bucket: {bucket_name}')
+        try:
+            return self.create_bucket(bucket_name, bucket_type, **kwargs)
+        except DuplicateBucketName:
+            self._duplicated_bucket_name_debug_info(bucket_name)
+            raise
+
+    @backoff.on_exception(
+        backoff.expo,
+        TooManyRequests,
+        max_tries=8,
+    )
+    def clean_bucket(self, bucket: Union[Bucket, str]) -> None:
+        if isinstance(bucket, str):
+            bucket = self.get_bucket_by_name(bucket)
+
+        files_leftover = False
+        file_versions = bucket.ls(latest_only=False, recursive=True)
+
+        for file_version_info, _ in file_versions:
+            if file_version_info.file_retention:
+                if file_version_info.file_retention.mode == RetentionMode.GOVERNANCE:
+                    print(f'Removing retention from file version: {file_version_info.id_}')
+                    self.update_file_retention(
+                        file_version_info.id_,
+                        file_version_info.file_name,
+                        NO_RETENTION_FILE_SETTING,
+                        bypass_governance=True
+                    )
+                elif file_version_info.file_retention.mode == RetentionMode.COMPLIANCE:
+                    if file_version_info.file_retention.retain_until > current_time_millis():  # yapf: disable
+                        print(
+                            f'File version: {file_version_info.id_} cannot be removed due to compliance mode retention'
+                        )
+                        files_leftover = True
+                        continue
+                elif file_version_info.file_retention.mode == RetentionMode.NONE:
+                    pass
+                else:
+                    raise ValueError(
+                        f'Unknown retention mode: {file_version_info.file_retention.mode}'
+                    )
+            if file_version_info.legal_hold.is_on():
+                print(f'Removing legal hold from file version: {file_version_info.id_}')
+                self.update_file_legal_hold(
+                    file_version_info.id_, file_version_info.file_name, LegalHold.OFF
+                )
+            print(f'Removing file version: {file_version_info.id_}')
+            try:
+                self.delete_file_version(file_version_info.id_, file_version_info.file_name)
+            except FileNotPresent:
+                print(
+                    f'It seems that file version {file_version_info.id_} has already been removed'
+                )
+
+        if files_leftover:
+            print('Unable to remove bucket because some retained files remain')
+        else:
+            print(f'Removing bucket: {bucket.name}')
+            try:
+                self.delete_bucket(bucket)
+            except (BucketIdNotFound, NonExistentBucket):
+                print(f'It seems that bucket {bucket.name} has already been removed')
+        print()
+
+    def clean_buckets(self) -> None:
+        self.count_and_print_buckets()
+        for bucket in self.buckets:
+            self.clean_bucket(bucket)
+        self.buckets = []
+
+    def clean_all_buckets(self) -> None:
+        buckets = self.list_buckets()
+        print(f'Total bucket count: {len(buckets)}')
+
+        for bucket in buckets:
+            if not bucket.name.startswith(BUCKET_NAME_PREFIX):
+                print(f'Skipping bucket removal: "{bucket.name}"')
+                continue
+            self.clean_bucket(bucket)
+
+        buckets = self.list_buckets()
+        print(f'Total bucket count after cleanup: {len(buckets)}')
+        for bucket in buckets:
+            print(bucket)
+
+    def count_and_print_buckets(self) -> None:
+        buckets = self.buckets
+        count = len(buckets)
+        print(f'Total bucket count at {datetime.now()}: {count}')
+        for i, bucket in enumerate(buckets, start=1):
+            print(f'- {i}\t{bucket.name} [{bucket.id_}]')
+
+    def _duplicated_bucket_name_debug_info(self, bucket_name: str) -> None:
+        # Trying to obtain as much information as possible about this bucket.
+        print(' DUPLICATED BUCKET DEBUG START '.center(60, '='))
+        bucket = self.get_bucket_by_name(bucket_name)
+
+        print('Bucket metadata:')
+        bucket_dict = bucket.as_dict()
+        for info_key, info in bucket_dict.items():
+            print('\t%s: "%s"' % (info_key, info))
+
+        print('All files (and their versions) inside the bucket:')
+        ls_generator = bucket.ls(recursive=True, latest_only=False)
+        for file_version, _directory in ls_generator:
+            # as_dict() is bound to have more info than we can use,
+            # but maybe some of it will cast some light on the issue.
+            print('\t%s (%s)' % (file_version.file_name, file_version.as_dict()))
+
+        print(' DUPLICATED BUCKET DEBUG END '.center(60, '='))

b2sdk/test/bucket_tracking.py

@@ -0,0 +1,32 @@
+######################################################################
+#
+# File: b2sdk/test/bucket_tracking.py
+#
+# Copyright 2023 Backblaze Inc. All Rights Reserved.
+#
+# License https://www.backblaze.com/using_b2_code.html
+#
+######################################################################
+from b2sdk.v2 import Bucket
+
+
+class BucketTrackingMixin:
+    """
+    Mixin class for B2Api, which enables bucket tracking.
+    This mixin will add a `buckets` member to the B2Api instance and will use it track created and
+    deleted buckets.  The main purpose of this are tests -- the `buckets` member can be used in test
+    teardown to ensure proper bucket cleanup.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.buckets = []
+        super().__init__(*args, **kwargs)
+
+    def create_bucket(self, name: str, *args, **kwargs) -> Bucket:
+        bucket = super().create_bucket(name, *args, **kwargs)
+        self.buckets.append(bucket)
+        return bucket
+
+    def delete_bucket(self, bucket: Bucket):
+        super().delete_bucket(bucket)
+        self.buckets = [b for b in self.buckets if b.id_ != bucket.id_]

noxfile.py

@@ -134,15 +134,15 @@ def unit(session):
 @nox.session(python=PYTHON_VERSIONS)
 def integration(session):
     """Run integration tests."""
-    install_myself(session)
+    install_myself(session, ['dev'])
     session.run('pip', 'install', *REQUIREMENTS_TEST)
     session.run('pytest', '-s', *session.posargs, 'test/integration')
 
 
 @nox.session(python=PYTHON_DEFAULT_VERSION)
 def cleanup_old_buckets(session):
     """Remove buckets from previous test runs."""
-    install_myself(session)
+    install_myself(session, ['dev'])
     session.run('pip', 'install', *REQUIREMENTS_TEST)
     session.run('python', '-m', 'test.integration.cleanup_buckets')
 

requirements-dev.txt

@@ -0,0 +1 @@
+backoff>=1.4.0,<3.0.0

setup.py

@@ -108,7 +108,10 @@ def read_requirements(extra=None):
     # dependencies). You can install these using the following syntax,
     # for example:
     # $ pip install -e .[dev,test]
-    extras_require={'doc': read_requirements('doc')},
+    extras_require={
+        'doc': read_requirements('doc'),
+        'dev': read_requirements('dev'),
+    },
     setup_requires=['setuptools_scm<6.0'],
     use_scm_version=True,
 

test/integration/__init__.py

@@ -8,9 +8,10 @@
 #
 ######################################################################
 import os
+from typing import Tuple
 
 
-def get_b2_auth_data():
+def get_b2_auth_data() -> Tuple[str, str]:
     application_key_id = os.environ.get('B2_TEST_APPLICATION_KEY_ID')
     if application_key_id is None:
         raise ValueError('B2_TEST_APPLICATION_KEY_ID is not set.')
@@ -19,3 +20,7 @@ def get_b2_auth_data():
     if application_key is None:
         raise ValueError('B2_TEST_APPLICATION_KEY is not set.')
     return application_key_id, application_key
+
+
+def get_realm() -> str:
+    return os.environ.get('B2_TEST_ENVIRONMENT', 'production')

test/integration/base.py

@@ -13,10 +13,7 @@
 
 import pytest
 
-from b2sdk.v2 import current_time_millis
-from b2sdk.v2.exception import DuplicateBucketName
-from .bucket_cleaner import BucketCleaner
-from .helpers import GENERAL_BUCKET_NAME_PREFIX, BUCKET_NAME_LENGTH, BUCKET_CREATED_AT_MILLIS, bucket_name_part, authorize
+from b2sdk.test.api_test_manager import ApiTestManager
 
 
 class IntegrationTestBase:
@@ -26,30 +23,11 @@ def set_http_debug(self):
             http.client.HTTPConnection.debuglevel = 1
 
     @pytest.fixture(autouse=True)
-    def save_settings(self, dont_cleanup_old_buckets, b2_auth_data):
-        type(self).dont_cleanup_old_buckets = dont_cleanup_old_buckets
-        type(self).b2_auth_data = b2_auth_data
-
-    @classmethod
-    def setup_class(cls):
-        cls.this_run_bucket_name_prefix = GENERAL_BUCKET_NAME_PREFIX + bucket_name_part(8)
-
-    @classmethod
-    def teardown_class(cls):
-        BucketCleaner(
-            cls.dont_cleanup_old_buckets,
-            *cls.b2_auth_data,
-            current_run_prefix=cls.this_run_bucket_name_prefix
-        ).cleanup_buckets()
-
-    @pytest.fixture(autouse=True)
-    def setup_method(self):
-        self.b2_api, self.info = authorize(self.b2_auth_data)
-
-    def generate_bucket_name(self):
-        return self.this_run_bucket_name_prefix + bucket_name_part(
-            BUCKET_NAME_LENGTH - len(self.this_run_bucket_name_prefix)
-        )
+    def setup_method(self, b2_auth_data, realm):
+        self.b2_api = ApiTestManager(*b2_auth_data, realm)
+        self.info = self.b2_api.account_info
+        yield
+        self.b2_api.clean_buckets()
 
     def write_zeros(self, file, number):
         line = b'0' * 1000 + b'\n'
@@ -60,32 +38,4 @@ def write_zeros(self, file, number):
             written += line_len
 
     def create_bucket(self):
-        bucket_name = self.generate_bucket_name()
-        try:
-            return self.b2_api.create_bucket(
-                bucket_name,
-                'allPublic',
-                bucket_info={BUCKET_CREATED_AT_MILLIS: str(current_time_millis())}
-            )
-        except DuplicateBucketName:
-            self._duplicated_bucket_name_debug_info(bucket_name)
-            raise
-
-    def _duplicated_bucket_name_debug_info(self, bucket_name: str) -> None:
-        # Trying to obtain as much information as possible about this bucket.
-        print(' DUPLICATED BUCKET DEBUG START '.center(60, '='))
-        bucket = self.b2_api.get_bucket_by_name(bucket_name)
-
-        print('Bucket metadata:')
-        bucket_dict = bucket.as_dict()
-        for info_key, info in bucket_dict.items():
-            print('\t%s: "%s"' % (info_key, info))
-
-        print('All files (and their versions) inside the bucket:')
-        ls_generator = bucket.ls(recursive=True, latest_only=False)
-        for file_version, _directory in ls_generator:
-            # as_dict() is bound to have more info than we can use,
-            # but maybe some of it will cast some light on the issue.
-            print('\t%s (%s)' % (file_version.file_name, file_version.as_dict()))
-
-        print(' DUPLICATED BUCKET DEBUG END '.center(60, '='))
+        return self.b2_api.create_test_bucket()