PR changes

data-dot-all · Sep 23, 2024 · 9bf30c3 · 9bf30c3
1 parent f46a6bd
commit 9bf30c3
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 25 deletions.
diff --git a/tests_new/integration_tests/aws_clients/athena.py b/tests_new/integration_tests/aws_clients/athena.py
@@ -16,7 +16,7 @@ def _run_query(self, query, workgroup='primary', output_location=None):
             result = self._client.start_query_execution(QueryString=query, WorkGroup=workgroup)
         return result['QueryExecutionId']
 
-    @poller(check_success=lambda state: state not in ['QUEUED', 'RUNNING'], timeout=5)
+    @poller(check_success=lambda state: state not in ['QUEUED', 'RUNNING'], timeout=600, sleep_time=5)
     def _wait_for_query(self, query_id):
         result = self._client.get_query_execution(QueryExecutionId=query_id)
         return result['QueryExecution']['Status']['State']
@@ -34,3 +34,4 @@ def get_env_work_group(self, env_name):
         for workgroup in workgroups:
             if env_name in workgroup:
                 return workgroup
+        return workgroups[0] if workgroups else None
diff --git a/tests_new/integration_tests/aws_clients/iam.py b/tests_new/integration_tests/aws_clients/iam.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 
@@ -9,9 +10,7 @@
 
 
 class IAMClient:
-    def __init__(self, session, region):
-        if session is None:
-            session = boto3.Session()
+    def __init__(self, session=boto3.Session(), region=os.environ.get('AWS_REGION', 'us-east-1')):
         self._client = session.client('iam', region_name=region)
         self._resource = session.resource('iam', region_name=region)
         self._region = region
@@ -29,29 +28,31 @@ def get_tooling_account_id():
         session = boto3.Session()
         param_client = session.client('ssm', os.environ.get('AWS_REGION', 'us-east-1'))
         parameter_path = f"/dataall/{os.environ.get('ENVNAME', 'dev')}/toolingAccount"
-        print(parameter_path)
         toolingAccount = param_client.get_parameter(Name=parameter_path)['Parameter']['Value']
         return toolingAccount
 
     def create_role(self, account_id, role_name, test_role_name):
+        policy_doc = {
+            'Version': '2012-10-17',
+            'Statement': [
+                {
+                    'Effect': 'Allow',
+                    'Principal': {
+                        'AWS': [
+                            f'arn:aws:iam::{account_id}:root',
+                            f'arn:aws:iam::{IAMClient.get_tooling_account_id()}:root',
+                            f'arn:aws:sts::{account_id}:assumed-role/{test_role_name}/{test_role_name}',
+                        ]
+                    },
+                    'Action': 'sts:AssumeRole',
+                    'Condition': {},
+                }
+            ],
+        }
         try:
             role = self._client.create_role(
                 RoleName=role_name,
-                AssumeRolePolicyDocument=f"""{{
-                                         "Version": "2012-10-17",
-            "Statement": [
-                {{
-                    "Effect": "Allow",
-                    "Principal": {{
-                        "AWS": ["arn:aws:iam::{account_id}:root",
-                        "arn:aws:iam::{IAMClient.get_tooling_account_id()}:root",
-                        "arn:aws:sts::{account_id}:assumed-role/{test_role_name}/{test_role_name}"]
-                    }},
-                    "Action": "sts:AssumeRole",
-                    "Condition": {{}}
-                }}
-            ]
-            }}""",
+                AssumeRolePolicyDocument=json.dumps(policy_doc),
                 Description='Role for Lambda function',
             )
             return role

diff --git a/tests_new/integration_tests/aws_clients/sts.py b/tests_new/integration_tests/aws_clients/sts.py
@@ -1,10 +1,10 @@
+import os
+
 import boto3
 
 
 class StsClient:
-    def __init__(self, session, region):
-        if not session:
-            session = boto3.Session()
+    def __init__(self, session=boto3.Session(), region=os.environ.get('AWS_REGION', 'us-east-1')):
         self._client = session.client('sts', region_name=region)
         self._region = region
 

diff --git a/tests_new/integration_tests/modules/share_base/conftest.py b/tests_new/integration_tests/modules/share_base/conftest.py
@@ -12,7 +12,7 @@
 from tests_new.integration_tests.modules.share_base.utils import check_share_ready
 from dataall.modules.shares_base.services.shares_enums import ShareItemStatus
 
-test_cons_role_name = 'dataall-test-ShareTestConsumptionRole'
+test_cons_role_name = 'dataall-test-ShareTestConsumptionRole2'
 
 
 def revoke_all_possible(client, shareUri):

diff --git a/tests_new/integration_tests/modules/share_base/test_new_crossacc_s3_share.py b/tests_new/integration_tests/modules/share_base/test_new_crossacc_s3_share.py
@@ -1,6 +1,7 @@
 import pytest
 from assertpy import assert_that
 
+from tests_new.integration_tests.modules.s3_datasets.queries import get_folder
 from tests_new.integration_tests.aws_clients.athena import AthenaClient
 from dataall.modules.shares_base.services.shares_enums import (
     ShareItemStatus,
@@ -250,7 +251,10 @@ def test_check_item_access(
             assert_that(s3_client.bucket_exists(item.itemName)).is_not_none()
             assert_that(s3_client.list_bucket_objects(item.itemName)).is_not_none()
         elif item.itemType == ShareableType.StorageLocation.name:
-            assert_that(s3_client.list_accesspoint_folder_objects(access_point_arn, item.itemName + '/')).is_not_none()
+            folder = get_folder(client5, item.itemUri)
+            assert_that(
+                s3_client.list_accesspoint_folder_objects(access_point_arn, folder.S3Prefix + '/')
+            ).is_not_none()
 
 
 @pytest.mark.dependency(name='share_revoked', depends=['share_succeeded'])