From 4819bc766040a0fd338edeb77573375f32bbca32 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Wed, 2 Oct 2024 08:46:42 +0200
Subject: [PATCH 1/9] clean-up dataset S3/KMS resources for created datasets

---
 .../modules/s3_datasets/global_conftest.py     | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
index 8de97095d..75c605697 100644
--- a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
+++ b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
@@ -95,7 +95,7 @@ def create_aws_imported_resources(
     return bucket, kms_alias, database, existing_lf_admins
 
 
-def delete_aws_imported_resources(aws_client, env, bucket=None, kms_alias=None, database=None, existing_lf_admins=None):
+def delete_aws_dataset_resources(aws_client, env, bucket=None, kms_alias=None, database=None, existing_lf_admins=None):
     try:
         if bucket:
             S3Client(session=aws_client, region=env['region']).delete_bucket(bucket)
@@ -218,7 +218,7 @@ def create_filters(client, tables):
 
 
 @pytest.fixture(scope='session')
-def session_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata):
+def session_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata, session_env1_aws_client):
     ds = None
     try:
         ds = create_s3_dataset(
@@ -235,6 +235,9 @@ def session_s3_dataset1(client1, group1, org1, session_env1, session_id, testdat
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
+            delete_aws_dataset_resources(
+                aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName, kms_alias=ds.KmsAlias
+            )
 
 
 @pytest.fixture(scope='session')
@@ -276,7 +279,7 @@ def session_imported_sse_s3_dataset1(
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
-        delete_aws_imported_resources(aws_client=session_env1_aws_client, env=session_env1, bucket=bucket)
+        delete_aws_dataset_resources(aws_client=session_env1_aws_client, env=session_env1, bucket=bucket)
 
 
 @pytest.fixture(scope='session')
@@ -331,7 +334,7 @@ def session_imported_kms_s3_dataset1(
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
-        delete_aws_imported_resources(
+        delete_aws_dataset_resources(
             aws_client=session_env1_aws_client,
             env=session_env1,
             bucket=bucket,
@@ -373,7 +376,7 @@ def session_imported_kms_s3_dataset1_tables_data_filters(client1, session_import
 
 
 @pytest.fixture(scope='function')
-def temp_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata):
+def temp_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata, session_env1_aws_client):
     ds = None
     try:
         ds = create_s3_dataset(
@@ -389,6 +392,9 @@ def temp_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata):
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
+            delete_aws_dataset_resources(
+                aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName, kms_alias=ds.KmsAlias
+            )
 
 
 """
@@ -495,7 +501,7 @@ def persistent_imported_kms_s3_dataset1(
         or (not kms_alias and not existing_kms_alias)
         or (not database and not existing_database)
     ):
-        delete_aws_imported_resources(
+        delete_aws_dataset_resources(
             aws_client=persistent_env1_aws_client,
             env=persistent_env1,
             bucket=bucket,

From 58999a7c5e1cc841d072103ff4e15a6df8556ef6 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Wed, 2 Oct 2024 10:21:10 +0200
Subject: [PATCH 2/9] clean-up environment bucket, skip KMS deletion

---
 .../core/environment/cdk/environment_stack.py |  9 +++++-
 tests_new/integration_tests/aws_clients/s3.py | 29 +++++++++++++++++++
 .../core/environment/global_conftest.py       |  4 +++
 .../modules/s3_datasets/global_conftest.py    |  9 ++----
 4 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 tests_new/integration_tests/aws_clients/s3.py

diff --git a/backend/dataall/core/environment/cdk/environment_stack.py b/backend/dataall/core/environment/cdk/environment_stack.py
index 451921ace..1518126fb 100644
--- a/backend/dataall/core/environment/cdk/environment_stack.py
+++ b/backend/dataall/core/environment/cdk/environment_stack.py
@@ -585,7 +585,14 @@ def create_integration_tests_role(self):
                     's3:DeleteObject',
                 ],
                 effect=iam.Effect.ALLOW,
-                resources=['arn:aws:s3:::dataalltesting*', 'arn:aws:s3:::dataalltesting*/*'],
+                resources=[
+                    'arn:aws:s3:::dataalltesting*',
+                    'arn:aws:s3:::dataalltesting*/*',
+                    'arn:aws:s3:::dataall-session*',
+                    'arn:aws:s3:::dataall-session*/*',
+                    'arn:aws:s3:::dataall-temp*',
+                    'arn:aws:s3:::dataall-temp*/*',
+                ],
             )
         )
         self.test_role.add_to_policy(
diff --git a/tests_new/integration_tests/aws_clients/s3.py b/tests_new/integration_tests/aws_clients/s3.py
new file mode 100644
index 000000000..5e960238e
--- /dev/null
+++ b/tests_new/integration_tests/aws_clients/s3.py
@@ -0,0 +1,29 @@
+import logging
+from botocore.exceptions import ClientError
+
+log = logging.getLogger(__name__)
+
+
+class S3Client:
+    def __init__(self, session, region):
+        self._client = session.client('s3', region_name=region)
+        self._resource = session.resource('s3', region_name=region)
+        self._region = region
+
+    def delete_bucket(self, bucket_name):
+        """
+        Delete an S3 bucket.
+        :param bucket_name: Name of the S3 bucket to be deleted
+        :return: None
+        """
+        try:
+            # Delete all objects in the bucket before deleting the bucket
+            bucket = self._resource.Bucket(bucket_name)
+            bucket_versioning = self._resource.BucketVersioning(bucket_name)
+            if bucket_versioning.status == 'Enabled':
+                bucket.object_versions.delete()
+            else:
+                bucket.objects.all().delete()
+            self._client.delete_bucket(Bucket=bucket_name)
+        except ClientError as e:
+            log.exception(f'Error deleting S3 bucket: {e}')
diff --git a/tests_new/integration_tests/core/environment/global_conftest.py b/tests_new/integration_tests/core/environment/global_conftest.py
index fcc1db311..005beb0c5 100644
--- a/tests_new/integration_tests/core/environment/global_conftest.py
+++ b/tests_new/integration_tests/core/environment/global_conftest.py
@@ -12,6 +12,7 @@
 )
 from integration_tests.core.organizations.queries import create_organization
 from integration_tests.core.stack.utils import check_stack_ready
+from tests_new.integration_tests.aws_clients.s3 import S3Client
 
 log = logging.getLogger(__name__)
 
@@ -50,6 +51,9 @@ def session_env1(client1, group1, org1, session_id, testdata):
         yield env
     finally:
         if env:
+            role = f'arn:aws:iam::{env.AwsAccountId}:role/dataall-integration-tests-role-{env.region}'
+            session = get_environment_aws_session(role, env)
+            S3Client(session=session, region=env.region).delete_bucket(env.EnvironmentDefaultBucketName)
             delete_env(client1, env)
 
 
diff --git a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
index 75c605697..4351e618b 100644
--- a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
+++ b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
@@ -235,9 +235,7 @@ def session_s3_dataset1(client1, group1, org1, session_env1, session_id, testdat
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
-            delete_aws_dataset_resources(
-                aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName, kms_alias=ds.KmsAlias
-            )
+            delete_aws_dataset_resources(aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName)
 
 
 @pytest.fixture(scope='session')
@@ -392,9 +390,8 @@ def temp_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata,
     finally:
         if ds:
             delete_s3_dataset(client1, session_env1['environmentUri'], ds)
-            delete_aws_dataset_resources(
-                aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName, kms_alias=ds.KmsAlias
-            )
+
+            delete_aws_dataset_resources(aws_client=session_env1_aws_client, env=session_env1, bucket=ds.S3BucketName)
 
 
 """

From 514441d66fc5ffea3ca6274e61a172550a7dd0b9 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Thu, 3 Oct 2024 16:45:54 +0200
Subject: [PATCH 3/9] Use S3 resource to delete objects in bucket

---
 tests_new/integration_tests/aws_clients/s3.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/tests_new/integration_tests/aws_clients/s3.py b/tests_new/integration_tests/aws_clients/s3.py
index 5e960238e..26ea4b512 100644
--- a/tests_new/integration_tests/aws_clients/s3.py
+++ b/tests_new/integration_tests/aws_clients/s3.py
@@ -19,11 +19,10 @@ def delete_bucket(self, bucket_name):
         try:
             # Delete all objects in the bucket before deleting the bucket
             bucket = self._resource.Bucket(bucket_name)
-            bucket_versioning = self._resource.BucketVersioning(bucket_name)
-            if bucket_versioning.status == 'Enabled':
-                bucket.object_versions.delete()
-            else:
-                bucket.objects.all().delete()
-            self._client.delete_bucket(Bucket=bucket_name)
+            # Delete all object versions
+            bucket.object_versions.all().delete()
+            # Delete any remaining objects (if versioning was not enabled)
+            bucket.objects.all().delete()
+            bucket.delete()
         except ClientError as e:
             log.exception(f'Error deleting S3 bucket: {e}')

From fbef997b554468661feac6c48c22e0d896aa05e6 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Sun, 6 Oct 2024 18:36:45 +0200
Subject: [PATCH 4/9] Fix permissions to delete S3 versions

---
 backend/dataall/core/environment/cdk/environment_stack.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/backend/dataall/core/environment/cdk/environment_stack.py b/backend/dataall/core/environment/cdk/environment_stack.py
index 1518126fb..9e195318e 100644
--- a/backend/dataall/core/environment/cdk/environment_stack.py
+++ b/backend/dataall/core/environment/cdk/environment_stack.py
@@ -583,6 +583,7 @@ def create_integration_tests_role(self):
                     's3:PutEncryptionConfiguration',
                     's3:GetObject*',
                     's3:DeleteObject',
+                    's3:DeleteObjectVersion',
                 ],
                 effect=iam.Effect.ALLOW,
                 resources=[

From f683d3ae7d532b1391899d45fc037132c47700b2 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Sun, 6 Oct 2024 18:51:52 +0200
Subject: [PATCH 5/9] Added delete access points

---
 .../dataall/core/environment/cdk/environment_stack.py    | 2 ++
 tests_new/integration_tests/aws_clients/s3.py            | 9 ++++++++-
 .../core/environment/global_conftest.py                  | 4 +++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/backend/dataall/core/environment/cdk/environment_stack.py b/backend/dataall/core/environment/cdk/environment_stack.py
index 9e195318e..51d5feac5 100644
--- a/backend/dataall/core/environment/cdk/environment_stack.py
+++ b/backend/dataall/core/environment/cdk/environment_stack.py
@@ -626,6 +626,8 @@ def create_integration_tests_role(self):
                     'kms:DescribeKey',
                     's3:GetBucketVersioning',
                     's3:List*',
+                    's3:ListAccessPoints',
+                    's3:DeleteAccessPoint',
                 ],
                 effect=iam.Effect.ALLOW,
                 resources=['*'],
diff --git a/tests_new/integration_tests/aws_clients/s3.py b/tests_new/integration_tests/aws_clients/s3.py
index 26ea4b512..6894e4f01 100644
--- a/tests_new/integration_tests/aws_clients/s3.py
+++ b/tests_new/integration_tests/aws_clients/s3.py
@@ -5,9 +5,11 @@
 
 
 class S3Client:
-    def __init__(self, session, region):
+    def __init__(self, session, account, region):
         self._client = session.client('s3', region_name=region)
+        self._control_client = session.client('s3control', region_name=region)
         self._resource = session.resource('s3', region_name=region)
+        self._account = account
         self._region = region
 
     def delete_bucket(self, bucket_name):
@@ -23,6 +25,11 @@ def delete_bucket(self, bucket_name):
             bucket.object_versions.all().delete()
             # Delete any remaining objects (if versioning was not enabled)
             bucket.objects.all().delete()
+            # Delete any remaining access point
+            paginator = self._control_client.get_paginator('list_access_points')
+            for page in paginator.paginate(AccountId=self._account, Bucket=bucket_name):
+                for access_point in page['AccessPointList']:
+                    self._control_client.delete_access_point(AccountId=self._account, Name=access_point['Name'])
             bucket.delete()
         except ClientError as e:
             log.exception(f'Error deleting S3 bucket: {e}')
diff --git a/tests_new/integration_tests/core/environment/global_conftest.py b/tests_new/integration_tests/core/environment/global_conftest.py
index 005beb0c5..1178181d9 100644
--- a/tests_new/integration_tests/core/environment/global_conftest.py
+++ b/tests_new/integration_tests/core/environment/global_conftest.py
@@ -53,7 +53,9 @@ def session_env1(client1, group1, org1, session_id, testdata):
         if env:
             role = f'arn:aws:iam::{env.AwsAccountId}:role/dataall-integration-tests-role-{env.region}'
             session = get_environment_aws_session(role, env)
-            S3Client(session=session, region=env.region).delete_bucket(env.EnvironmentDefaultBucketName)
+            S3Client(session=session, account=env.account, region=env.region).delete_bucket(
+                env.EnvironmentDefaultBucketName
+            )
             delete_env(client1, env)
 
 

From 6d1b98082ca89777a448f96f5d08b287860f19c7 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Mon, 7 Oct 2024 15:19:01 +0200
Subject: [PATCH 6/9] remove error from exception

---
 tests_new/integration_tests/modules/s3_datasets/aws_clients.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
index 020d74d2d..9b6716217 100644
--- a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
+++ b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
@@ -81,7 +81,7 @@ def delete_bucket(self, bucket_name):
                 bucket.objects.all().delete()
             self._client.delete_bucket(Bucket=bucket_name)
         except ClientError as e:
-            log.exception(f'Error deleting S3 bucket: {e}')
+            log.exception('Error deleting S3 bucket')
 
     def upload_file_to_prefix(self, local_file_path, s3_path):
         """

From 172ad10181d8f2410938de4afe2e179c64a877af Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Mon, 7 Oct 2024 15:20:43 +0200
Subject: [PATCH 7/9] remove error from exception other functions

---
 .../modules/s3_datasets/aws_clients.py        | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

diff --git a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
index 9b6716217..67d7f56e3 100644
--- a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
+++ b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
@@ -63,7 +63,7 @@ def create_bucket(self, bucket_name, kms_key_arn=None):
             )
             return bucket_name
         except ClientError as e:
-            log.exception(f'Error creating S3 bucket: {e}')
+            log.exception('Error creating S3 bucket')
 
     def delete_bucket(self, bucket_name):
         """
@@ -147,7 +147,7 @@ def create_key_with_alias(self, alias_name):
             return key_id, alias_name
 
         except ClientError as e:
-            log.exception(f'Error creating KMS key with alias: {e}')
+            log.exception('Error creating KMS key with alias')
 
     def _put_key_policy(self, key_id):
         response = self._client.get_key_policy(KeyId=key_id, PolicyName='default')
@@ -178,11 +178,7 @@ def _put_key_policy(self, key_id):
         try:
             self._client.put_key_policy(KeyId=key_id, PolicyName='default', Policy=json.dumps(policy))
         except ClientError as err:
-            log.exception(
-                "Couldn't set policy for key %s. Here's why %s",
-                key_id,
-                err,
-            )
+            log.exception("Couldn't set policy for key.")
 
     def delete_key_by_alias(self, alias_name):
         """
@@ -197,7 +193,7 @@ def delete_key_by_alias(self, alias_name):
                 self._client.schedule_key_deletion(KeyId=key_id)
             self._client.delete_alias(AliasName=f'alias/{alias_name}')
         except ClientError as e:
-            log.exception(f'Error deleting KMS key by alias: {e}')
+            log.exception('Error deleting KMS key by alias')
 
     def _get_key_by_alias(self, alias_name):
         try:
@@ -211,7 +207,7 @@ def _get_key_by_alias(self, alias_name):
             return None
 
         except ClientError as e:
-            log.exception(f'Error getting KMS key by alias: {e}')
+            log.exception('Error getting KMS key by alias')
 
 
 class GlueClient:
@@ -228,7 +224,7 @@ def get_database(self, database_name):
             database = self._client.get_database(Name=database_name)
             return database
         except ClientError as e:
-            log.exception(f'Database not found, exception: {e}')
+            log.exception('Database not found')
             return False
 
     def create_database(self, database_name, bucket):
@@ -237,7 +233,7 @@ def create_database(self, database_name, bucket):
             self._client.create_database(DatabaseInput={'Name': database_name, 'LocationUri': f's3://{bucket}/'})
             return database_name
         except ClientError as e:
-            log.exception(f'Error creating Glue database: {e}')
+            log.exception('Error creating Glue database')
 
     def create_table(self, database_name, bucket, table_name):
         try:
@@ -265,7 +261,7 @@ def create_table(self, database_name, bucket, table_name):
                 },
             )
         except ClientError as e:
-            log.exception(f'Error creating Glue table: {e}')
+            log.exception('Error creating Glue table')
 
     def delete_database(self, database_name):
         """
@@ -279,7 +275,7 @@ def delete_database(self, database_name):
             if e.response['Error']['Code'] == 'EntityNotFoundException':
                 log.exception(f"Glue database '{database_name}' does not exist.")
             else:
-                log.exception(f'Error deleting Glue database: {e}')
+                log.exception('Error deleting Glue database')
 
 
 class LakeFormationClient:
@@ -304,7 +300,7 @@ def add_role_to_datalake_admin(self, role_arn):
             )
             return existing_admins
         except ClientError as e:
-            log.exception(f'Error granting lake formation permissions: {e}')
+            log.exception('Error granting lake formation permissions')
 
     def remove_role_from_datalake_admin(self, old_existing_principals):
         try:
@@ -317,7 +313,7 @@ def remove_role_from_datalake_admin(self, old_existing_principals):
             )
             return True
         except ClientError as e:
-            log.exception(f'Error granting lake formation permissions: {e}')
+            log.exception('Error granting lake formation permissions')
 
     def grant_create_database(self, role_arn):
         """
@@ -333,4 +329,4 @@ def grant_create_database(self, role_arn):
             )
             return True
         except ClientError as e:
-            log.exception(f'Error granting permissions to create database: {e}')
+            log.exception('Error granting permissions to create database')

From e014991357d4acc6dfcf307b7223518a435d2b22 Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Tue, 8 Oct 2024 11:01:19 +0200
Subject: [PATCH 8/9] Simplify delete functions + no paginator list_access
 points

---
 tests_new/integration_tests/aws_clients/s3.py  |  9 +++++----
 .../core/environment/global_conftest.py        |  2 +-
 .../modules/s3_datasets/aws_clients.py         | 18 ------------------
 .../modules/s3_datasets/global_conftest.py     |  4 ++--
 4 files changed, 8 insertions(+), 25 deletions(-)

diff --git a/tests_new/integration_tests/aws_clients/s3.py b/tests_new/integration_tests/aws_clients/s3.py
index 6894e4f01..080105edb 100644
--- a/tests_new/integration_tests/aws_clients/s3.py
+++ b/tests_new/integration_tests/aws_clients/s3.py
@@ -26,10 +26,11 @@ def delete_bucket(self, bucket_name):
             # Delete any remaining objects (if versioning was not enabled)
             bucket.objects.all().delete()
             # Delete any remaining access point
-            paginator = self._control_client.get_paginator('list_access_points')
-            for page in paginator.paginate(AccountId=self._account, Bucket=bucket_name):
-                for access_point in page['AccessPointList']:
-                    self._control_client.delete_access_point(AccountId=self._account, Name=access_point['Name'])
+            access_points = self._control_client.list_access_points(AccountId=self._account, Bucket=bucket_name)[
+                'AccessPointList'
+            ]
+            for access_point in access_points:
+                self._control_client.delete_access_point(AccountId=self._account, Name=access_point['Name'])
             bucket.delete()
         except ClientError as e:
             log.exception(f'Error deleting S3 bucket: {e}')
diff --git a/tests_new/integration_tests/core/environment/global_conftest.py b/tests_new/integration_tests/core/environment/global_conftest.py
index 46f33ba0c..176a9a5fb 100644
--- a/tests_new/integration_tests/core/environment/global_conftest.py
+++ b/tests_new/integration_tests/core/environment/global_conftest.py
@@ -56,7 +56,7 @@ def session_env1(client1, group1, org1, session_id, testdata):
         if env:
             role = f'arn:aws:iam::{env.AwsAccountId}:role/dataall-integration-tests-role-{env.region}'
             session = get_environment_aws_session(role, env)
-            S3Client(session=session, account=env.account, region=env.region).delete_bucket(
+            S3Client(session=session, account=env.AwsAccountId, region=env.region).delete_bucket(
                 env.EnvironmentDefaultBucketName
             )
             delete_env(client1, env)
diff --git a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
index 67d7f56e3..4435bb4cf 100644
--- a/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
+++ b/tests_new/integration_tests/modules/s3_datasets/aws_clients.py
@@ -65,24 +65,6 @@ def create_bucket(self, bucket_name, kms_key_arn=None):
         except ClientError as e:
             log.exception('Error creating S3 bucket')
 
-    def delete_bucket(self, bucket_name):
-        """
-        Delete an S3 bucket.
-        :param bucket_name: Name of the S3 bucket to be deleted
-        :return: None
-        """
-        try:
-            # Delete all objects in the bucket before deleting the bucket
-            bucket = self._resource.Bucket(bucket_name)
-            bucket_versioning = self._resource.BucketVersioning(bucket_name)
-            if bucket_versioning.status == 'Enabled':
-                bucket.object_versions.delete()
-            else:
-                bucket.objects.all().delete()
-            self._client.delete_bucket(Bucket=bucket_name)
-        except ClientError as e:
-            log.exception('Error deleting S3 bucket')
-
     def upload_file_to_prefix(self, local_file_path, s3_path):
         """
         Upload a file from a local path to an S3 bucket with a specified prefix.
diff --git a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
index 9b79a3ff1..cbb36c845 100644
--- a/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
+++ b/tests_new/integration_tests/modules/s3_datasets/global_conftest.py
@@ -19,7 +19,7 @@
     create_table_data_filter,
 )
 from tests_new.integration_tests.modules.datasets_base.queries import list_datasets
-
+from integration_tests.aws_clients.s3 import S3Client as S3CommonClient
 from integration_tests.modules.s3_datasets.aws_clients import S3Client, KMSClient, GlueClient, LakeFormationClient
 from integration_tests.core.stack.queries import update_stack
 
@@ -101,7 +101,7 @@ def create_aws_imported_resources(
 def delete_aws_dataset_resources(aws_client, env, bucket=None, kms_alias=None, database=None, existing_lf_admins=None):
     try:
         if bucket:
-            S3Client(session=aws_client, region=env['region']).delete_bucket(bucket)
+            S3CommonClient(session=aws_client, account=env.AwsAccountId, region=env.region).delete_bucket(bucket)
         if kms_alias:
             KMSClient(
                 session=aws_client,

From b9b5f007737182b031d175ed626788fa3d3c3c9b Mon Sep 17 00:00:00 2001
From: dlpzx <dlpzx@amazon.com>
Date: Tue, 8 Oct 2024 17:53:33 +0200
Subject: [PATCH 9/9] Replace get_environment_aws_session by STS refreshable
 client

---
 tests_new/integration_tests/core/environment/global_conftest.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests_new/integration_tests/core/environment/global_conftest.py b/tests_new/integration_tests/core/environment/global_conftest.py
index 59293a024..b8f32ceec 100644
--- a/tests_new/integration_tests/core/environment/global_conftest.py
+++ b/tests_new/integration_tests/core/environment/global_conftest.py
@@ -55,7 +55,7 @@ def session_env1(client1, group1, org1, session_id, testdata):
     finally:
         if env:
             role = f'arn:aws:iam::{env.AwsAccountId}:role/dataall-integration-tests-role-{env.region}'
-            session = get_environment_aws_session(role, env)
+            session = STSClient(role_arn=role, region=env.region, session_name='Session_1').get_refreshable_session()
             S3Client(session=session, account=env.AwsAccountId, region=env.region).delete_bucket(
                 env.EnvironmentDefaultBucketName
             )