CICD Integration tests: s3 dataset shares, persistent shares (#1580)

### Feature or Bugfix
- Feature


### Detail
- test for persistent shares
- common test functions are places in separate file
- updated envs/dataset fixtures

### Relates
- #1376

### Security
Please answer the questions below briefly where applicable, or write
`N/A`. Based on
[OWASP 10](https://owasp.org/Top10/en/).

- Does this PR introduce or modify any input fields or queries - this
includes
fetching data from storage outside the application (e.g. a database, an
S3 bucket)?
  - Is the input sanitized?
- What precautions are you taking before deserializing the data you
consume?
  - Is injection prevented by parametrizing queries?
  - Have you ensured no `eval` or similar functions are used?
- Does this PR introduce any functionality or component that requires
authorization?
- How have you ensured it respects the existing AuthN/AuthZ mechanisms?
  - Are you logging failed auth attempts?
- Are you using or adding any cryptographic features?
  - Do you use a standard proven implementations?
  - Are the used keys controlled by the customer? Where are they stored?
- Are you introducing any new policies/roles/users?
  - Have you used the least-privilege principle? How?


By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

---------

Co-authored-by: dlpzx <[email protected]>
Co-authored-by: Noah Paige <[email protected]>
Co-authored-by: Sofia Sazonova <[email protected]>

backend/dataall/core/environment/db/environment_repositories.py

@@ -310,3 +310,16 @@ def query_all_active_environments(session) -> List[Environment]:
     @staticmethod
     def query_environment_groups(session, uri):
         return session.query(EnvironmentGroup).filter(EnvironmentGroup.environmentUri == uri).all()
+
+    @staticmethod
+    def get_environment_consumption_role_by_name(session, uri, IAMRoleName):
+        return (
+            session.query(ConsumptionRole)
+            .filter(
+                and_(
+                    ConsumptionRole.environmentUri == uri,
+                    ConsumptionRole.IAMRoleName == IAMRoleName,
+                )
+            )
+            .first()
+        )

backend/dataall/core/environment/services/environment_service.py

@@ -1147,3 +1147,9 @@ def resolve_consumption_role_policies(uri, IAMRoleName):
             region=environment.region,
             resource_prefix=environment.resourcePrefix,
         ).get_all_policies()
+
+    @staticmethod
+    @ResourcePolicyService.has_resource_permission(environment_permissions.GET_ENVIRONMENT)
+    def get_consumption_role_by_name(uri, IAMRoleName):
+        with get_context().db_engine.scoped_session() as session:
+            return EnvironmentRepository.get_environment_consumption_role_by_name(session, uri, IAMRoleName)

backend/dataall/core/stacks/api/types.py

@@ -21,6 +21,7 @@
         gql.Field(name='region', type=gql.NonNullableType(gql.String)),
         gql.Field(name='status', type=gql.String),
         gql.Field(name='stackid', type=gql.String),
+        gql.Field(name='updated', type=gql.AWSDateTime),
         gql.Field(name='link', type=gql.String, resolver=resolve_link),
         gql.Field(name='outputs', type=gql.String, resolver=resolve_outputs),
         gql.Field(name='resources', type=gql.String, resolver=resolve_resources),

tests_new/integration_tests/aws_clients/iam.py

@@ -16,19 +16,11 @@ def __init__(self, session=boto3.Session(), region=os.environ.get('AWS_REGION',
         self._region = region
 
     def get_role(self, role_name):
-        try:
-            role = self._client.get_role(RoleName=role_name)
-            return role
-        except Exception as e:
-            log.info(f'Error occurred: {e}')
-            return None
+        role = self._client.get_role(RoleName=role_name)
+        return role
 
     def delete_role(self, role_name):
-        try:
-            self._client.delete_role(RoleName=role_name)
-        except Exception as e:
-            log.error(e)
-            raise e
+        self._client.delete_role(RoleName=role_name)
 
     def create_role(self, account_id, role_name, test_role_name):
         policy_doc = {
@@ -47,16 +39,12 @@ def create_role(self, account_id, role_name, test_role_name):
                 }
             ],
         }
-        try:
-            role = self._client.create_role(
-                RoleName=role_name,
-                AssumeRolePolicyDocument=json.dumps(policy_doc),
-                Description='Role for Lambda function',
-            )
-            return role
-        except Exception as e:
-            log.error(e)
-            raise e
+        role = self._client.create_role(
+            RoleName=role_name,
+            AssumeRolePolicyDocument=json.dumps(policy_doc),
+            Description='Role for Lambda function',
+        )
+        return role
 
     def get_consumption_role(self, account_id, role_name, test_role_name):
         role = self.get_role(role_name)

tests_new/integration_tests/aws_clients/utils.py

@@ -0,0 +1,18 @@
+import json
+import boto3
+
+from tests_new.integration_tests.aws_clients.sts import StsClient
+
+
+def get_group_session(credentials_str):
+    credentials = json.loads(credentials_str)
+    return boto3.Session(
+        aws_access_key_id=credentials['AccessKey'],
+        aws_secret_access_key=credentials['SessionKey'],
+        aws_session_token=credentials['sessionToken'],
+    )
+
+
+def get_role_session(session, role_arn, region):
+    sts_client = StsClient(session=session, region=region)
+    return sts_client.get_role_session(role_arn)

tests_new/integration_tests/core/environment/global_conftest.py

@@ -1,4 +1,6 @@
 import logging
+from datetime import datetime
+
 import pytest
 import boto3
 
@@ -12,6 +14,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.core.environment.utils import update_env_stack
 
 log = logging.getLogger(__name__)
 
@@ -178,3 +181,30 @@ def get_or_create_persistent_env(env_name, client, group, testdata):
 @pytest.fixture(scope='session')
 def persistent_env1(client1, group1, testdata):
     return get_or_create_persistent_env('persistent_env1', client1, group1, testdata)
+
+
+@pytest.fixture(scope='session')
+def updated_persistent_env1(client1, group1, persistent_env1):
+    update_env_stack(client1, persistent_env1)
+    return get_environment(client1, persistent_env1.environmentUri)
+
+
+@pytest.fixture(scope='session')
+def persistent_cross_acc_env_1(client5, group5, testdata):
+    return get_or_create_persistent_env('persistent_cross_acc_env_1', client5, group5, testdata)
+
+
+@pytest.fixture(scope='session')
+def updated_persistent_cross_acc_env_1(client5, group5, persistent_cross_acc_env_1):
+    update_env_stack(client5, persistent_cross_acc_env_1)
+    return get_environment(client5, persistent_cross_acc_env_1.environmentUri)
+
+
+@pytest.fixture(scope='session')
+def persistent_cross_acc_env_1_integration_role_arn(persistent_cross_acc_env_1):
+    return f'arn:aws:iam::{persistent_cross_acc_env_1.AwsAccountId}:role/dataall-integration-tests-role-{persistent_cross_acc_env_1.region}'
+
+
+@pytest.fixture(scope='session')
+def persistent_cross_acc_env_1_aws_client(persistent_cross_acc_env_1, persistent_cross_acc_env_1_integration_role_arn):
+    return get_environment_aws_session(persistent_cross_acc_env_1_integration_role_arn, persistent_cross_acc_env_1)

tests_new/integration_tests/core/environment/queries.py

@@ -33,6 +33,7 @@
   accountid
   region
   stackid
+  updated
   link
   outputs
   resources
@@ -244,6 +245,33 @@ def add_consumption_role(client, env_uri, group_uri, consumption_role_name, iam_
     return response.data.addConsumptionRoleToEnvironment
 
 
+def list_environment_consumption_roles(client, env_uri, filter):
+    query = {
+        'operationName': 'listEnvironmentConsumptionRoles',
+        'variables': {'environmentUri': env_uri, 'filter': filter},
+        'query': """
+                    query listEnvironmentConsumptionRoles($environmentUri: String!, $filter: ConsumptionRoleFilter) {
+                      listEnvironmentConsumptionRoles(environmentUri: $environmentUri, filter: $filter) {
+                        count
+                        page
+                        pages
+                        hasNext
+                        hasPrevious
+                        nodes {
+                          consumptionRoleUri
+                          consumptionRoleName
+                          environmentUri
+                          groupUri
+                          IAMRoleArn
+                        }
+                      }
+                    }
+        """,
+    }
+    response = client.query(query=query)
+    return response.data.listEnvironmentConsumptionRoles
+
+
 def remove_consumption_role(client, env_uri, consumption_role_uri):
     query = {
         'operationName': 'removeConsumptionRoleFromEnvironment',

tests_new/integration_tests/core/environment/test_environment.py

@@ -12,10 +12,13 @@
     remove_consumption_role,
     remove_group_from_env,
 )
-from integration_tests.core.stack.queries import update_stack
-from integration_tests.core.stack.utils import check_stack_in_progress, check_stack_ready
+
 from integration_tests.errors import GqlError
 
+from integration_tests.core.environment.utils import update_env_stack
+from integration_tests.core.stack.queries import get_stack
+
+
 log = logging.getLogger(__name__)
 
 
@@ -51,15 +54,18 @@ def test_list_envs_invited(client2, session_env1, session_env2, session_id):
 
 
 def test_persistent_env_update(client1, persistent_env1):
-    # wait for stack to get to a final state before triggering an update
-    stack_uri = persistent_env1.stack.stackUri
-    env_uri = persistent_env1.environmentUri
-    check_stack_ready(client1, env_uri, stack_uri)
-    update_stack(client1, env_uri, 'environment')
-    # wait for stack to move to "in_progress" state
-    check_stack_in_progress(client1, env_uri, stack_uri)
-    stack = check_stack_ready(client1, env_uri, stack_uri)
-    assert_that(stack.status).is_equal_to('UPDATE_COMPLETE')
+    stack = get_stack(
+        client1,
+        persistent_env1.environmentUri,
+        persistent_env1.stack.stackUri,
+        persistent_env1.environmentUri,
+        target_type='environment',
+    )
+    updated_before = datetime.fromisoformat(stack.updated)
+    stack = update_env_stack(client1, persistent_env1)
+    assert_that(stack).contains_entry(status='UPDATE_COMPLETE')
+    updated = datetime.fromisoformat(stack.updated)
+    assert_that(updated).is_greater_than_or_equal_to(updated_before)
 
 
 def test_invite_group_on_env_no_org(client1, session_env2, group4):

tests_new/integration_tests/core/environment/utils.py

@@ -1,5 +1,6 @@
 from integration_tests.core.environment.queries import update_environment
 from integration_tests.core.stack.utils import check_stack_ready, check_stack_in_progress
+from integration_tests.core.stack.queries import update_stack
 
 
 def set_env_params(client, env, **new_params):
@@ -34,3 +35,14 @@ def set_env_params(client, env, **new_params):
         )
         check_stack_in_progress(client, env_uri, stack_uri)
         check_stack_ready(client, env_uri, stack_uri)
+
+
+def update_env_stack(client, env):
+    stack_uri = env.stack.stackUri
+    env_uri = env.environmentUri
+    # wait for stack to get to a final state before triggering an update
+    check_stack_ready(client, env_uri, stack_uri)
+    update_stack(client, env_uri, 'environment')
+    # wait for stack to move to "in_progress" state
+    check_stack_in_progress(client, env_uri, stack_uri)
+    return check_stack_ready(client, env_uri, stack_uri)

tests_new/integration_tests/core/stack/queries.py

@@ -44,6 +44,7 @@ def get_stack(client, env_uri, stack_uri, target_uri, target_type):
                         accountid
                         region
                         stackid
+                        updated
                         link
                         outputs
                         resources

tests_new/integration_tests/modules/s3_datasets/aws_clients.py

@@ -104,7 +104,7 @@ def list_bucket_objects(self, bucket_name):
             return self._client.list_objects(Bucket=bucket_name)
         except ClientError as e:
             logging.error(f'Error listing objects in S3: {e}')
-            raise
+            raise e
 
     def list_accesspoint_folder_objects(self, access_point, folder_name):
         try:

tests_new/integration_tests/modules/s3_datasets/global_conftest.py

@@ -1,4 +1,6 @@
 import logging
+import time
+
 import pytest
 import boto3
 import json
@@ -19,6 +21,7 @@
 from tests_new.integration_tests.modules.datasets_base.queries import list_datasets
 
 from integration_tests.modules.s3_datasets.aws_clients import S3Client, KMSClient, GlueClient, LakeFormationClient
+from integration_tests.core.stack.queries import update_stack
 
 log = logging.getLogger(__name__)
 
@@ -398,7 +401,15 @@ def temp_s3_dataset1(client1, group1, org1, session_env1, session_id, testdata):
 
 
 def get_or_create_persistent_s3_dataset(
-    dataset_name, client, group, env, autoApprovalEnabled=False, bucket=None, kms_alias='', glue_database=''
+    dataset_name,
+    client,
+    group,
+    env,
+    autoApprovalEnabled=False,
+    bucket=None,
+    kms_alias='',
+    glue_database='',
+    withContent=False,
 ):
     dataset_name = dataset_name or 'persistent_s3_dataset1'
     s3_datasets = list_datasets(client, term=dataset_name).nodes
@@ -431,6 +442,9 @@ def get_or_create_persistent_s3_dataset(
                 tags=[dataset_name],
                 autoApprovalEnabled=autoApprovalEnabled,
             )
+            if withContent:
+                create_tables(client, s3_dataset)
+                create_folders(client, s3_dataset)
 
         if s3_dataset.stack.status in ['CREATE_COMPLETE', 'UPDATE_COMPLETE']:
             return s3_dataset
@@ -441,7 +455,21 @@ def get_or_create_persistent_s3_dataset(
 
 @pytest.fixture(scope='session')
 def persistent_s3_dataset1(client1, group1, persistent_env1, testdata):
-    return get_or_create_persistent_s3_dataset('persistent_s3_dataset1', client1, group1, persistent_env1)
+    return get_or_create_persistent_s3_dataset(
+        'persistent_s3_dataset1', client1, group1, persistent_env1, withContent=True
+    )
+
+
+@pytest.fixture(scope='session')
+def updated_persistent_s3_dataset1(client1, persistent_s3_dataset1):
+    target_type = 'dataset'
+    stack_uri = persistent_s3_dataset1.stack.stackUri
+    env_uri = persistent_s3_dataset1.environment.environmentUri
+    dataset_uri = persistent_s3_dataset1.datasetUri
+    update_stack(client1, dataset_uri, target_type)
+    time.sleep(120)
+    check_stack_ready(client1, env_uri=env_uri, stack_uri=stack_uri, target_uri=dataset_uri, target_type=target_type)
+    return get_dataset(client1, dataset_uri)
 
 
 @pytest.fixture(scope='session')