diff --git a/CHANGELOG.md b/CHANGELOG.md index d9d402768..7d14de18e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +* Add `BucketTrackingMixin` class for tracking buckets created by `B2Api` instances + ## [1.18.0] - 2022-09-20 ### Added diff --git a/b2sdk/_v3/__init__.py b/b2sdk/_v3/__init__.py index 26583f542..49e86e0f7 100644 --- a/b2sdk/_v3/__init__.py +++ b/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 --git a/b2sdk/api.py b/b2sdk/api.py index c38866049..1934931b5 100644 --- a/b2sdk/api.py +++ b/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 --git a/test/unit/v_all/test_api.py b/test/unit/v_all/test_api.py index 6e3dcfeee..60c2c2aed 100644 --- a/test/unit/v_all/test_api.py +++ b/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_}