Added redshift-connection tests

data-dot-all · Oct 10, 2024 · b21c8da · b21c8da
1 parent d367e3e
commit b21c8da
Show file tree

Hide file tree

Showing 4 changed files with 361 additions and 34 deletions.
diff --git a/backend/dataall/modules/redshift_datasets/aws/redshift.py b/backend/dataall/modules/redshift_datasets/aws/redshift.py
@@ -14,7 +14,14 @@ def __init__(self, account_id: str, region: str) -> None:
 
     def describe_cluster(self, clusterId: str):
         log.info(f'Describing cluster {clusterId=}')
-        return self.client.describe_clusters(ClusterIdentifier=clusterId)['Clusters'][0]
+        try:
+            return self.client.describe_clusters(ClusterIdentifier=clusterId)['Clusters'][0]
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'ClusterNotFound':
+                log.error(f'Redshift cluster {clusterId} does not exist')
+                return None
+            else:
+                raise e
 
     def get_cluster_namespaceId(self, clusterId: str):
         log.info(f'Describing cluster {clusterId=}')

diff --git a/tests_new/integration_tests/modules/redshift_datasets/connection_queries.py b/tests_new/integration_tests/modules/redshift_datasets/connection_queries.py
@@ -128,16 +128,16 @@ def list_redshift_connection_group_no_permissions(client, connection_uri, term='
 def create_redshift_connection(
     client,
     connection_name,
-    connection_type,
     environment_uri,
     group_uri,
     redshift_type,
-    cluster_id,
-    namespace_id,
-    workgroup,
     database,
-    redshift_user,
-    secret_arn,
+    connection_type,
+    cluster_id=None,
+    namespace_id=None,
+    workgroup=None,
+    redshift_user=None,
+    secret_arn=None,
 ):
     query = {
         'operationName': 'createRedshiftConnection',

diff --git a/tests_new/integration_tests/modules/redshift_datasets/global_conftest.py b/tests_new/integration_tests/modules/redshift_datasets/global_conftest.py
@@ -6,6 +6,8 @@
 from integration_tests.modules.redshift_datasets.connection_queries import (
     create_redshift_connection,
     delete_redshift_connection,
+    add_redshift_connection_group_permissions,
+    delete_redshift_connection_group_permissions,
 )
 
 log = logging.getLogger(__name__)
@@ -64,6 +66,24 @@ def session_connection_serverless_admin(client1, group1, session_env1, testdata)
             delete_redshift_connection(client=client1, connection_uri=connection.connectionUri)
 
 
+@pytest.fixture(scope='session')
+def session_connection_serverless_admin_group_with_permissions(client1, group2, session_connection_serverless_admin):
+    permissions = None
+    try:
+        permissions = add_redshift_connection_group_permissions(
+            client=client1,
+            connection_uri=session_connection_serverless_admin.connectionUri,
+            group_uri=group2,
+            permissions=['CREATE_SHARE_REQUEST_WITH_CONNECTION'],
+        )
+        yield group2
+    finally:
+        if permissions:
+            delete_redshift_connection_group_permissions(
+                client=client1, connection_uri=session_connection_serverless_admin.connectionUri, group_uri=group2
+            )
+
+
 @pytest.fixture(scope='session')
 def session_connection_serverless_data_user(client1, group1, session_env1, testdata):
     connection = None
@@ -84,41 +104,41 @@ def session_connection_serverless_data_user(client1, group1, session_env1, testd
 
 
 @pytest.fixture(scope='session')
-def session_connection_cluster_admin(client1, group1, session_cross_acc_env_1, testdata):
+def session_connection_cluster_admin(client5, group5, session_cross_acc_env_1, testdata):
     connection = None
     try:
         connection = create_connection(
-            client=client1,
-            name='connection_serverless_admin_session_env1',
+            client=client5,
+            name='connection_cluster_admin_session_cross_acc_env_1',
             conn_type='ADMIN',
             env=session_cross_acc_env_1,
-            group=group1,
+            group=group5,
             red_type='cluster',
-            connection_data=testdata.redshift_connections['connection_serverless_admin_session_env1'],
+            connection_data=testdata.redshift_connections['connection_cluster_admin_session_cross_acc_env_1'],
         )
         yield connection
     finally:
         if connection:
-            delete_redshift_connection(client=client1, connection_uri=connection.connectionUri)
+            delete_redshift_connection(client=client5, connection_uri=connection.connectionUri)
 
 
 @pytest.fixture(scope='session')
-def session_connection_cluster_data_user(client1, group1, session_cross_acc_env_1, testdata):
+def session_connection_cluster_data_user(client5, group5, session_cross_acc_env_1, testdata):
     connection = None
     try:
         connection = create_connection(
-            client=client1,
-            name='connection_serverless_data_user_session_env1',
+            client=client5,
+            name='connection_cluster_data_user_session_cross_acc_env_1',
             conn_type='DATA_USER',
             env=session_cross_acc_env_1,
-            group=group1,
+            group=group5,
             red_type='cluster',
-            connection_data=testdata.redshift_connections['connection_serverless_data_user_session_env1'],
+            connection_data=testdata.redshift_connections['connection_cluster_data_user_session_cross_acc_env_1'],
         )
         yield connection
     finally:
         if connection:
-            delete_redshift_connection(client=client1, connection_uri=connection.connectionUri)
+            delete_redshift_connection(client=client5, connection_uri=connection.connectionUri)
 
 
 @pytest.fixture(scope='session')