2025-04-03 10:39:47 +05:30

87 lines
3.0 KiB
Python

# Copyright 2025 Collate
# Licensed under the Collate Community License, Version 1.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Test Redshift using the topology
"""
from unittest import TestCase
from unittest.mock import patch
from metadata.generated.schema.metadataIngestion.workflow import (
OpenMetadataWorkflowConfig,
)
from metadata.ingestion.source.database.redshift.metadata import RedshiftSource
mock_redshift_config = {
"source": {
"type": "redshift",
"serviceName": "local_redshift",
"serviceConnection": {
"config": {
"type": "Redshift",
"username": "username",
"password": "password",
"database": "database",
"hostPort": "cluster.name.region.redshift.amazonaws.com:5439",
"sslMode": "verify-full",
"sslConfig": {
"caCertificate": "CA certificate content",
},
}
},
"sourceConfig": {"config": {"type": "DatabaseMetadata"}},
},
"sink": {"type": "metadata-rest", "config": {}},
"workflowConfig": {
"openMetadataServerConfig": {
"hostPort": "http://localhost:8585/api",
"authProvider": "openmetadata",
"securityConfig": {"jwtToken": "redshift"},
}
},
}
RAW_DIST_STYLE = ["KEY(eventid)", "EVEN", "ALL"]
EXPECTED_PARTITION_COLUMNS = ["eventid", None, None]
class RedshiftUnitTest(TestCase):
@patch(
"metadata.ingestion.source.database.common_db_source.CommonDbSourceService.test_connection"
)
def __init__(self, methodName, test_connection) -> None:
super().__init__(methodName)
test_connection.return_value = False
self.config = OpenMetadataWorkflowConfig.model_validate(mock_redshift_config)
self.redshift_source = RedshiftSource.create(
mock_redshift_config["source"],
self.config.workflowConfig.openMetadataServerConfig,
)
def test_partition_parse_columns(self):
for i in range(len(RAW_DIST_STYLE)):
with self.subTest(i=i):
self.assertEqual(
self.redshift_source._get_partition_key(RAW_DIST_STYLE[i]),
EXPECTED_PARTITION_COLUMNS[i],
)
@patch("sqlalchemy.engine.base.Engine")
@patch(
"metadata.ingestion.source.database.common_db_source.CommonDbSourceService.connection"
)
def test_close_connection(self, engine, connection):
connection.return_value = True
self.redshift_source.close()