OpenMetadata/ingestion/tests/unit/helpers_test.py

#  Licensed to the Apache Software Foundation (ASF) under one or more
#  contributor license agreements. See the NOTICE file distributed with
#  this work for additional information regarding copyright ownership.
#  The ASF licenses this file to You under the Apache License, Version 2.0
#  (the "License"); you may not use this file except in compliance with
#  the License. You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  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.
import json
from unittest import TestCase

from metadata.generated.schema.api.data.createDatabase import (
    CreateDatabaseEntityRequest,
)
from metadata.generated.schema.api.data.createTable import CreateTableEntityRequest
from metadata.generated.schema.api.services.createDashboardService import (
    CreateDashboardServiceEntityRequest,
)
from metadata.generated.schema.api.services.createDatabaseService import (
    CreateDatabaseServiceEntityRequest,
)
from metadata.generated.schema.api.services.createMessagingService import (
    CreateMessagingServiceEntityRequest,
)
from metadata.generated.schema.entity.data.table import Column
from metadata.generated.schema.type.entityReference import EntityReference
from metadata.ingestion.ometa.client import APIError
from metadata.ingestion.ometa.openmetadata_rest import (
    MetadataServerConfig,
    OpenMetadataAPIClient,
)


class RestTest(TestCase):
    file_path = "tests/unit/mysql_test.json"
    with open(file_path) as ingestionFile:
        ingestionData = ingestionFile.read()
    client_config = json.loads(ingestionData).get("metadata_server")
    config = client_config.get("config", {})
    metadata_config = MetadataServerConfig.parse_obj(config)
    openmetadata_client = OpenMetadataAPIClient(metadata_config)
    client = OpenMetadataAPIClient(metadata_config).client

    def test_1_create_service(self):
        data = {
            "jdbc": {
                "connectionUrl": "mysql://localhost/openmetadata_db",
                "driverClass": "jdbc",
            },
            "name": "local_mysql_test",
            "serviceType": "MySQL",
            "description": "local mysql env",
        }
        create_mysql_service = CreateDatabaseServiceEntityRequest(**data)
        mysql_service = self.openmetadata_client.create_database_service(
            create_mysql_service
        )
        self.assertEqual(mysql_service.name, create_mysql_service.name)

    def test_2_get_service(self):
        mysql_service = self.openmetadata_client.get_database_service(
            "local_mysql_test"
        )
        self.assertEqual(mysql_service.name, "local_mysql_test")

    def test_3_get_service_by_id(self):
        mysql_service = self.openmetadata_client.get_database_service(
            "local_mysql_test"
        )
        mysql_service_get_id = self.openmetadata_client.get_database_service_by_id(
            mysql_service.id.__root__
        )
        self.assertEqual(mysql_service.id, mysql_service_get_id.id)

    def test_4_create_update_databases(self):
        mysql_service = self.openmetadata_client.get_database_service(
            "local_mysql_test"
        )
        service_reference = EntityReference(
            id=mysql_service.id.__root__, type="databaseService"
        )
        create_database_request = CreateDatabaseEntityRequest(
            name="dwh", service=service_reference
        )
        created_database = self.openmetadata_client.create_database(
            create_database_request
        )
        created_database.description = "hello world"
        update_database_request = CreateDatabaseEntityRequest(
            name=created_database.name,
            description=created_database.description,
            service=service_reference,
        )
        updated_database = self.openmetadata_client.create_database(
            update_database_request
        )
        self.assertEqual(updated_database.description, created_database.description)

    def test_5_create_table(self):
        databases = self.openmetadata_client.list_databases()
        columns = [
            Column(name="id", columnDataType="INT"),
            Column(name="name", columnDataType="VARCHAR"),
        ]
        table = CreateTableEntityRequest(
            name="test1", columns=columns, database=databases[0].id.__root__
        )
        created_table = self.openmetadata_client.create_or_update_table(table)
        self.client.delete(f"/tables/{created_table.id.__root__}")
        self.client.delete(f"/databases/{databases[0].id.__root__}")
        self.assertEqual(table.name, created_table.name)

    def test_6_delete_service(self):
        mysql_service = self.openmetadata_client.get_database_service(
            "local_mysql_test"
        )
        self.openmetadata_client.delete_database_service(mysql_service.id.__root__)
        self.assertRaises(
            APIError,
            self.openmetadata_client.get_database_service_by_id,
            mysql_service.id.__root__,
        )

    def test_7_create_messaging_service(self):
        create_messaging_service = CreateMessagingServiceEntityRequest(
            name="sample_kafka_test",
            serviceType="Kafka",
            brokers=["localhost:9092"],
            schemaRegistry="http://localhost:8081",
        )
        messaging_service = self.openmetadata_client.create_messaging_service(
            create_messaging_service
        )
        self.assertEqual(create_messaging_service.name, messaging_service.name)

    def test_8_get_messaging_service(self):
        messaging_service = self.openmetadata_client.get_messaging_service(
            "sample_kafka_test"
        )
        self.client.delete(
            f"/services/messagingServices/{messaging_service.id.__root__}"
        )
        self.assertEqual(messaging_service.name, "sample_kafka_test")

    def test_9_create_dashboard_service(self):
        create_dashboard_service = CreateDashboardServiceEntityRequest(
            name="sample_superset_test",
            serviceType="Superset",
            username="admin",
            password="admin",
            dashboardUrl="http://localhost:8088",
        )
        dashboard_service = None
        try:
            dashboard_service = self.openmetadata_client.create_dashboard_service(
                create_dashboard_service
            )
        except APIError:
            print(APIError)
        self.assertEqual(create_dashboard_service.name, dashboard_service.name)

    def test_10_get_dashboard_service(self):
        dashboard_service = self.openmetadata_client.get_dashboard_service(
            "sample_superset_test"
        )
        self.client.delete(
            f"/services/dashboardServices/{dashboard_service.id.__root__}"
        )
        self.assertEqual(dashboard_service.name, "sample_superset_test")