Integration Tests Notebooks

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

@@ -621,3 +621,34 @@ def create_integration_tests_role(self):
                 resources=['*'],
             )
         )
+
+        self.test_role.add_to_policy(
+            iam.PolicyStatement(
+                actions=['ec2:Describe*', 'ec2:*Vpc*', 'ec2:*Subnet*', 'ec2:*Route*', 'ec2:*Tags*'],
+                effect=iam.Effect.ALLOW,
+                resources=[
+                    'arn:aws:ec2:*:139956106467:route-table/*',
+                    'arn:aws:ec2:*:139956106467:security-group/*',
+                    'arn:aws:ec2:*:139956106467:vpc/*',
+                    'arn:aws:ec2:*:139956106467:security-group-rule/*',
+                    'arn:aws:ec2:*:139956106467:subnet/*',
+                    'arn:aws:ec2:*:139956106467:network-acl/*',
+                ],
+            ),
+        )
+
+        self.test_role.add_to_policy(
+            iam.PolicyStatement(
+                actions=['ec2:Describe*'],
+                effect=iam.Effect.ALLOW,
+                resources=['*'],
+            ),
+        )
+
+        self.test_role.add_to_policy(
+            iam.PolicyStatement(
+                actions=['cloudformation:Describe*'],
+                effect=iam.Effect.ALLOW,
+                resources=['arn:aws:cloudformation:*:139956106467:stack/*/*'],
+            ),
+        )

backend/dataall/modules/notebooks/db/notebook_repository.py

@@ -46,10 +46,12 @@ def _query_user_notebooks(self, username, groups, filter) -> Query:
             )
         )
         if filter and filter.get('term'):
+            term = filter['term']
             query = query.filter(
                 or_(
-                    SagemakerNotebook.description.ilike(filter.get('term') + '%%'),
-                    SagemakerNotebook.label.ilike(filter.get('term') + '%%'),
+                    SagemakerNotebook.description.ilike(term + '%%'),
+                    SagemakerNotebook.label.ilike(term + '%%'),
+                    SagemakerNotebook.tags.contains(f'{{{term}}}'),
                 )
             )
         return query.order_by(SagemakerNotebook.label)

deploy/stacks/param_store_stack.py

@@ -125,7 +125,7 @@ def __init__(
             f'toolingAccountParam{envname}',
             parameter_name=f'/dataall/{envname}/toolingAccount',
             string_value=str(tooling_account_id),
-            description=f'Stores AWS account if for the tooling account that hosts the code for environment {envname}',
+            description=f'Store AWS account if for the tooling account that hosts the code for environment {envname}',
         )
 
         if prod_sizing:

tests_new/integration_tests/client.py

@@ -2,7 +2,7 @@
 import boto3
 import os
 from munch import DefaultMunch
-
+from retrying import retry
 from integration_tests.errors import GqlError
 
 ENVNAME = os.getenv('ENVNAME', 'dev')
@@ -14,6 +14,17 @@ def __init__(self, username, password):
         self.password = password
         self.token = self._get_jwt_token()
 
+    @staticmethod
+    def _retry_if_connection_error(exception):
+        """Return True if we should retry, False otherwise"""
+        return isinstance(exception, requests.exceptions.ConnectionError) or isinstance(exception, requests.ReadTimeout)
+
+    @retry(
+        retry_on_exception=_retry_if_connection_error,
+        stop_max_attempt_number=3,
+        wait_random_min=1000,
+        wait_random_max=3000,
+    )
     def query(self, query: str):
         graphql_endpoint = os.path.join(os.environ['API_ENDPOINT'], 'graphql', 'api')
         headers = {'AccessKeyId': 'none', 'SecretKey': 'none', 'authorization': self.token}

tests_new/integration_tests/core/environment/global_conftest.py

@@ -51,19 +51,13 @@ def session_env1(client1, group1, org1, session_id, testdata):
             delete_env(client1, env)
 
 
-@pytest.fixture(scope='session')
-def session_env1_integration_role_arn(session_env1):
-    yield f'arn:aws:iam::{session_env1.AwsAccountId}:role/dataall-integration-tests-role-{session_env1.region}'
-
-
-@pytest.fixture(scope='session')
-def session_env1_aws_client(session_env1, session_id, session_env1_integration_role_arn):
+def get_environment_aws_session(role_arn, env):
     try:
         base_session = boto3.Session()
-        response = base_session.client('sts', region_name=session_env1.region).assume_role(
-            RoleArn=session_env1_integration_role_arn, RoleSessionName=session_env1_integration_role_arn.split('/')[1]
+        response = base_session.client('sts', region_name=env.region).assume_role(
+            RoleArn=role_arn, RoleSessionName=role_arn.split('/')[1]
         )
-        yield boto3.Session(
+        return boto3.Session(
             aws_access_key_id=response['Credentials']['AccessKeyId'],
             aws_secret_access_key=response['Credentials']['SecretAccessKey'],
             aws_session_token=response['Credentials']['SessionToken'],
@@ -73,29 +67,24 @@ def session_env1_aws_client(session_env1, session_id, session_env1_integration_r
         raise
 
 
+@pytest.fixture(scope='session')
+def session_env1_integration_role_arn(session_env1):
+    return f'arn:aws:iam::{session_env1.AwsAccountId}:role/dataall-integration-tests-role-{session_env1.region}'
+
+
+@pytest.fixture(scope='session')
+def session_env1_aws_client(session_env1, session_env1_integration_role_arn):
+    return get_environment_aws_session(session_env1_integration_role_arn, session_env1)
+
+
 @pytest.fixture(scope='session')
 def persistent_env1_integration_role_arn(persistent_env1):
-    yield f'arn:aws:iam::{persistent_env1.AwsAccountId}:role/dataall-integration-tests-role-{persistent_env1.region}'
+    return f'arn:aws:iam::{persistent_env1.AwsAccountId}:role/dataall-integration-tests-role-{persistent_env1.region}'
 
 
 @pytest.fixture(scope='session')
-def persistent_env1_aws_client(persistent_env1, session_id):
-    try:
-        base_session = boto3.Session()
-        role_arn = (
-            f'arn:aws:iam::{persistent_env1.AwsAccountId}:role/dataall-integration-tests-role-{persistent_env1.region}'
-        )
-        response = base_session.client('sts', region_name=persistent_env1.region).assume_role(
-            RoleArn=role_arn, RoleSessionName=role_arn.split('/')[1]
-        )
-        yield boto3.Session(
-            aws_access_key_id=response['Credentials']['AccessKeyId'],
-            aws_secret_access_key=response['Credentials']['SecretAccessKey'],
-            aws_session_token=response['Credentials']['SessionToken'],
-        )
-    except:
-        log.exception('Failed to assume environment integration test role')
-        raise
+def persistent_env1_aws_client(persistent_env1, persistent_env1_integration_role_arn):
+    return get_environment_aws_session(persistent_env1_integration_role_arn, persistent_env1)
 
 
 @pytest.fixture(scope='session')

tests_new/integration_tests/core/environment/queries.py

@@ -61,6 +61,9 @@ def create_environment(client, name, group, organizationUri, awsAccountId, regio
                 'region': region,
                 'description': 'Created for integration testing',
                 'tags': tags,
+                'parameters': [
+                    {'key': 'notebooksEnabled', 'value': 'true'},
+                ],
                 'type': 'IntegrationTesting',
             }
         },

tests_new/integration_tests/core/stack/utils.py

@@ -16,3 +16,15 @@ def check_stack_in_progress(client, env_uri, stack_uri, target_uri=None, target_
 @poller(check_success=lambda stack: not is_stack_in_progress(stack), timeout=600)
 def check_stack_ready(client, env_uri, stack_uri, target_uri=None, target_type='environment'):
     return get_stack(client, env_uri, stack_uri, target_uri or env_uri, target_type)
+
+
+def wait_stack_delete_complete(cf_client, stack_name):
+    # Wait for the stack to be deleted
+    waiter = cf_client.get_waiter('stack_delete_complete')
+    waiter.wait(
+        StackName=stack_name,
+        WaiterConfig={
+            'Delay': 20,  # Delay between each poll request (in seconds)
+            'MaxAttempts': 60,  # Maximum number of attempts before giving up
+        },
+    )

tests_new/integration_tests/modules/notebooks/aws_clients.py

@@ -0,0 +1,97 @@
+import logging
+import json
+from typing import Any, Dict, Optional
+from botocore.exceptions import ClientError
+
+log = logging.getLogger(__name__)
+
+
+class VpcClient:
+    def __init__(self, session, region):
+        self._client = session.client('ec2', region_name=region)
+        self._region = region
+
+    def create_vpc(self, vpc_name: str, cidr: str = '10.0.0.0/28') -> str:
+        log.info('Creating VPC..')
+        response = self._client.create_vpc(
+            CidrBlock=cidr,
+            TagSpecifications=[
+                {
+                    'ResourceType': 'vpc',
+                    'Tags': [
+                        {'Key': 'Name', 'Value': vpc_name},
+                    ],
+                },
+            ],
+        )
+        vpc_id = response['Vpc']['VpcId']
+        log.info(f'VPC created with ID: {vpc_id=}')
+        return vpc_id
+
+    def delete_vpc(self, vpc_id: str) -> Dict[str, Any]:
+        log.info('Deleting VPC..')
+        response = self._client.delete_vpc(VpcId=vpc_id)
+        log.info(f'VPC deleted with ID: {vpc_id=}')
+        return response
+
+    def get_vpc_id_by_name(self, vpc_name: str) -> Optional[str]:
+        log.info('Getting VPC ID by name..')
+        response = self._client.describe_vpcs(Filters=[{'Name': 'tag:Name', 'Values': [vpc_name]}])
+        if len(response['Vpcs']) == 0:
+            log.info(f'VPC with name {vpc_name} not found')
+            return None
+        vpc_id = response['Vpcs'][0]['VpcId']
+        log.info(f'VPC ID found: {vpc_id=}')
+        return vpc_id
+
+    def delete_vpc_by_name(self, vpc_name: str):
+        try:
+            vpc_id = self.get_vpc_id_by_name(vpc_name)
+            if vpc_id:
+                self.delete_vpc(vpc_id)
+            return True
+        except Exception as e:
+            log.error(f'Error deleting vpc {vpc_name=}. Error Message: {e}')
+
+    def create_subnet(self, vpc_id: str, subnet_name: str, cidr: str) -> str:
+        log.info('Creating subnet..')
+        response = self._client.create_subnet(
+            VpcId=vpc_id,
+            CidrBlock=cidr,
+            TagSpecifications=[
+                {
+                    'ResourceType': 'subnet',
+                    'Tags': [
+                        {'Key': 'Name', 'Value': subnet_name},
+                    ],
+                },
+            ],
+        )
+        subnet_id = response['Subnet']['SubnetId']
+        log.info(f'Subnet created with ID: {subnet_id=}')
+        return subnet_id
+
+    def get_subnet_id_by_name(self, subnet_name: str) -> Optional[str]:
+        log.info('Getting subnet ID by name..')
+        response = self._client.describe_subnets(Filters=[{'Name': 'tag:Name', 'Values': [subnet_name]}])
+        if len(response['Subnets']) == 0:
+            log.info(f'Subnet with name {subnet_name} not found')
+            return None
+        subnet_id = response['Subnets'][0]['SubnetId']
+        log.info(f'Subnet ID found: {subnet_id=}')
+        return subnet_id
+
+    def delete_subnet(self, subnet_id: str) -> Dict[str, Any]:
+        log.info('Deleting subnet..')
+        response = self._client.delete_subnet(SubnetId=subnet_id)
+        log.info(f'Subnet deleted with ID: {subnet_id=}')
+        return response
+
+    def delete_subnet_by_name(self, subnet_name: str):
+        try:
+            subnet_id = self.get_subnet_id_by_name(subnet_name)
+            if subnet_id:
+                self.delete_subnet(subnet_id)
+            return True
+        except Exception as e:
+            log.error(f'Error deleting subnet {subnet_name=}. Error Message: {e}')