Add BucketTrackingMixin class for tracking buckets created by B2Api instances

Author: mlesniewski-reef

Diff for: CHANGELOG.md

@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Logging performance summary of parallel download threads
 * Add `max_download_streams_per_file` parameter to B2Api class and underlying structures
 * Add `is_file_lock_enabled` parameter to `Bucket.update()` and related methods
+* Add `BucketTrackingMixin` class for tracking buckets created by `B2Api` instances
 
 ### Fixed
 * Replace `ReplicationScanResult.source_has_sse_c_enabled` with `source_encryption_mode`

Diff for: b2sdk/_v3/__init__.py

@@ -13,7 +13,7 @@
 
 # core
 
-from b2sdk.api import B2Api
+from b2sdk.api import B2Api, BucketTrackingMixin
 from b2sdk.api import Services
 from b2sdk.bucket import Bucket
 from b2sdk.bucket import BucketFactory

Diff for: b2sdk/api.py

@@ -595,3 +595,26 @@ def _check_bucket_restrictions(self, key, value):
         if allowed_bucket_identifier is not None:
             if allowed_bucket_identifier != value:
                 raise RestrictedBucket(allowed_bucket_identifier)
+
+
+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, *args, **kwargs):
+        bucket = super().create_bucket(name, *args, **kwargs)
+        self.buckets.append(bucket)
+        return bucket
+
+    def delete_bucket(self, bucket):
+        super().delete_bucket(bucket)
+        self.buckets = [b for b in self.buckets if b.id_ != bucket.id_]

Diff for: test/unit/v_all/test_api.py

@@ -10,7 +10,7 @@
 
 import pytest
 
-from apiver_deps import B2Api
+from apiver_deps import B2Api, BucketTrackingMixin
 from apiver_deps import B2HttpApiConfig
 from apiver_deps import Bucket
 from apiver_deps import InMemoryCache
@@ -102,11 +102,13 @@ def test_api_initialization(self, kwargs, _raw_api_class):
         assert download_manager.strategies[0].max_streams == kwargs['max_download_streams_per_file']
 
 
-class TestApi(TestBase):
+class TestApiBase(TestBase):
+    B2_API_CLASS = B2Api
+
     def setUp(self):
         self.account_info = InMemoryAccountInfo()
         self.cache = InMemoryCache()
-        self.api = B2Api(
+        self.api = self.B2_API_CLASS(
             self.account_info, self.cache, api_config=B2HttpApiConfig(_raw_api_class=RawSimulator)
         )
         self.raw_api = self.api.session.raw_api
@@ -115,6 +117,8 @@ def setUp(self):
     def _authorize_account(self):
         self.api.authorize_account('production', self.application_key_id, self.master_key)
 
+
+class TestApi(TestApiBase):
     @pytest.mark.apiver(to_ver=1)
     def test_get_bucket_by_id_up_to_v1(self):
         bucket = self.api.get_bucket_by_id("this id doesn't even exist")
@@ -158,3 +162,24 @@ def test_get_download_url_for_fileid(self):
         download_url = self.api.get_download_url_for_fileid('file-id')
 
         assert download_url == 'http://download.example.com/b2api/v2/b2_download_file_by_id?fileId=file-id'
+
+
+class TestBucketTrackingMixin(TestApiBase):
+    class BucketTrackingApi(BucketTrackingMixin, B2Api):
+        pass
+
+    B2_API_CLASS = BucketTrackingApi
+
+    def test_bucket_tracking(self):
+        self._authorize_account()
+
+        bucket_1, bucket_2, bucket_3 = [
+            self.api.create_bucket(f'bucket-{i + 1}', 'allPrivate') for i in range(3)
+        ]
+
+        self.api.delete_bucket(bucket_2)
+        self.api.delete_bucket(bucket_3)
+
+        bucket_4 = self.api.create_bucket('bucket-4', 'allPrivate')
+
+        assert {bucket.id_ for bucket in self.api.buckets} == {bucket_1.id_, bucket_4.id_}