CICD Integration tests: s3 dataset shares, persistent shares

backend/dataall/core/environment/api/queries.py

@@ -22,6 +22,7 @@
     list_groups,
     list_valid_environments,
     get_consumption_role_policies,
+    get_consumption_role,
 )
 from dataall.core.environment.api.types import (
     Environment,
@@ -220,3 +221,15 @@
     resolver=get_consumption_role_policies,
     test_scope='Environment',
 )
+
+
+getConsumptionRole = gql.QueryField(
+    name='getConsumptionRole',
+    args=[
+        gql.Argument(name='environmentUri', type=gql.NonNullableType(gql.String)),
+        gql.Argument(name='IAMRoleName', type=gql.NonNullableType(gql.String)),
+    ],
+    type=gql.Ref('ConsumptionRole'),
+    resolver=get_consumption_role,
+    test_scope='Environment',
+)

backend/dataall/core/environment/api/resolvers.py

@@ -285,3 +285,7 @@ def resolve_parameters(context, source: Environment, **kwargs):
     if not source:
         return None
     return EnvironmentService.get_environment_parameters(source.environmentUri)
+
+
+def get_consumption_role(context, source, environmentUri, IAMRoleName):
+    return EnvironmentService.get_consumption_role_by_name(uri=environmentUri, IAMRoleName=IAMRoleName)

backend/dataall/core/environment/cdk/environment_stack.py

@@ -657,7 +657,7 @@ def create_integration_tests_role(self):
 
         self.test_role.add_to_policy(
             iam.PolicyStatement(
-                actions=['iam:GetRole', 'iam:CreateRole', 'iam:PutRolePolicy'],
+                actions=['iam:GetRole', 'iam:CreateRole', 'iam:DeleteRole', 'iam:PutRolePolicy'],
                 effect=iam.Effect.ALLOW,
                 resources=[f'arn:aws:iam::{self.account}:role/dataall-test-*'],
             ),

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

@@ -71,6 +71,13 @@ def get_consumption_role(self, account_id, role_name, test_role_name):
             self.put_consumption_role_policy(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
+
     def put_consumption_role_policy(self, role_name):
         self._client.put_role_policy(
             RoleName=role_name,

tests_new/integration_tests/aws_clients/utils.py

@@ -0,0 +1,43 @@
+import json
+import boto3
+
+from tests_new.integration_tests.aws_clients.sts import StsClient
+
+
+# it's here and not in Env test module, because it's used only here and we don't want circular dependencies
+def get_environment_access_token(client, env_uri, group_uri):
+    query = {
+        'operationName': 'generateEnvironmentAccessToken',
+        'variables': {
+            'environmentUri': env_uri,
+            'groupUri': group_uri,
+        },
+        'query': """
+                     query generateEnvironmentAccessToken(
+                      $environmentUri: String!
+                      $groupUri: String
+                    ) {
+                      generateEnvironmentAccessToken(
+                        environmentUri: $environmentUri
+                        groupUri: $groupUri
+                      )
+                    }
+        """,
+    }
+    response = client.query(query=query)
+    return response.data.generateEnvironmentAccessToken
+
+
+def get_group_session(client, env_uri, group):
+    credentials = json.loads(get_environment_access_token(client, env_uri, group))
+
+    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/conftest.py

@@ -167,6 +167,11 @@ def session_id() -> str:
     return datetime.datetime.utcnow().isoformat()
 
 
+@pytest.fixture(scope='session', autouse=True)
+def session_start_timestamp(session_id) -> float:
+    yield datetime.datetime.fromisoformat(session_id).timestamp()
+
+
 @pytest.fixture(scope='session')
 def resources_prefix(session_id) -> str:
     re.sub('[^a-zA-Z0-9-]', '', session_id).lower()

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,35 @@ def add_consumption_role(client, env_uri, group_uri, consumption_role_name, iam_
     return response.data.addConsumptionRoleToEnvironment
 
 
+def get_consumption_role(client, env_uri, iam_role_name):
+    query = {
+        'operationName': 'getConsumptionRole',
+        'variables': {
+            'environmentUri': env_uri,
+            'IAMRoleName': iam_role_name,
+        },
+        'query': """
+                    query getConsumptionRole(
+                      $environmentUri: String!
+                      $IAMRoleName: String!
+                    ) {
+                      getConsumptionRole(
+                        environmentUri: $environmentUri
+                        IAMRoleName: $IAMRoleName
+                      ) {
+                        consumptionRoleUri
+                        consumptionRoleName
+                        environmentUri
+                        groupUri
+                        IAMRoleArn
+                      }
+                    }
+        """,
+    }
+    response = client.query(query=query)
+    return response.data.getConsumptionRole
+
+
 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.strptime(stack.updated, '%Y-%m-%d %H:%M:%S.%f').timestamp()
+    stack = update_env_stack(client1, persistent_env1)
+    assert_that(stack).contains_entry(status='UPDATE_COMPLETE')
+    updated = datetime.strptime(stack.updated, '%Y-%m-%d %H:%M:%S.%f').timestamp()
+    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')