OpenMetadata/ingestion/tests/cli_e2e/test_cli_tableau.py

#  Copyright 2022 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 Tableau connector with CLI - Enhanced with comprehensive lineage and metadata testing
"""

from pathlib import Path
from typing import List  # noqa: UP035

import pytest

from metadata.generated.schema.entity.data.chart import Chart
from metadata.generated.schema.entity.data.dashboard import Dashboard
from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
from metadata.ingestion.api.status import Status

from .base.test_cli import PATH_TO_RESOURCES  # noqa: TID252
from .common.test_cli_dashboard import CliCommonDashboard  # noqa: TID252


class TableauExpectedValues:
    """
    Centralized expected values for Tableau testing based on actual test data
    """

    # Basic counts - Updated based on actual test knowledge base
    DASHBOARDS_AND_CHARTS = 5  # 1 dashboard + 4 charts = 5
    LINEAGE_EDGES = 3
    TAGS = 2
    DATAMODEL_LINEAGE = 0
    DATAMODELS = 2
    FILTERED_MIX = 2
    FILTERED_SINK_MIX = 2
    DASHBOARDS_AND_CHARTS_AFTER_PATCH = 5
    SERVICE_NAME = "local_tableau"

    # Expected entity names
    EXPECTED_DASHBOARD_NAMES = ["Analytics Workbook"]  # noqa: RUF012

    EXPECTED_CHART_NAMES = [  # noqa: RUF012
        "Product Measure Sheet",
        "Sales Story",
        "Product vs Category Dashboard",
        "Category Measure Sheet",
    ]

    EXPECTED_DATAMODEL_NAMES = [  # noqa: RUF012
        "Sales Summary"  # Appears in both TableauEmbeddedDatasource and TableauPublishedDatasource
    ]

    # Expected data model columns/fields
    EXPECTED_DATAMODEL_FIELDS = [  # noqa: RUF012
        "state",
        "category_name",
        "order_date",
        "product_id",
        "customer_id",
        "line_total",
        "region",
        "product_name",
        "customer_name",
        "Sales Summary (Custom SQL)",
        "price_at_purchase",
        "order_id",
        "quantity",
    ]

    # Expected tags
    EXPECTED_TAGS = ["Analytics", "workbook"]  # noqa: RUF012

    # Expected chart type
    EXPECTED_CHART_TYPE = "ChartType.Other"

    # Expected field type
    EXPECTED_FIELD_TYPE = "tableau field"

    # Expected data model types
    EXPECTED_DATAMODEL_TYPES = [  # noqa: RUF012
        "DataModelType.TableauEmbeddedDatasource",
        "DataModelType.TableauPublishedDatasource",
    ]

    # Expected SQL for data models
    EXPECTED_DATAMODEL_SQL = """SELECT
        o.order_id,
        o.order_date,
        c.customer_id,
        c.first_name || ' ' || c.last_name AS customer_name,
        c.state,
        c.region,
        p.product_id,
        p.product_name,
        cat.category_name,
        oi.quantity,
        oi.price_at_purchase,
        (oi.quantity * oi.price_at_purchase) AS line_total
    FROM
        sales.orders AS o
    JOIN
        sales.customers AS c ON o.customer_id = c.customer_id
    JOIN
        sales.order_items AS oi ON o.order_id = oi.order_id
    JOIN
        inventory.products AS p ON oi.product_id = p.product_id
    JOIN
        inventory.categories AS cat ON p.category_id = cat.category_id"""

    # Lineage expectations: Tables -> TableauPublishedDatasource -> TableauEmbeddedDatasource -> Dashboard
    EXPECTED_SOURCE_TABLES = [  # noqa: RUF012
        "categories",
        "customers",
        "order_items",
        "orders",
        "products",
    ]

    # Filter patterns
    INCLUDE_DASHBOARDS = [".*Analytics.*"]  # noqa: RUF012
    EXCLUDE_DASHBOARDS = ["Sample.*"]  # noqa: RUF012
    INCLUDE_CHARTS = [".*Sheet.*", ".*Product.*", ".*Sales.*"]  # noqa: RUF012
    EXCLUDE_CHARTS = ["Obesity"]  # noqa: RUF012
    INCLUDE_DATAMODELS = [".*Sales.*", ".*Summary.*"]  # noqa: RUF012
    EXCLUDE_DATAMODELS = ["Random.*"]  # noqa: RUF012


class TableauCliTest(CliCommonDashboard.TestSuite):
    """
    Enhanced Tableau CLI Test Suite with comprehensive lineage and metadata validation
    """

    def prepare(self) -> None:
        """Prepare test environment by setting up required database service"""
        redshift_file_path = str(Path(PATH_TO_RESOURCES + f"/dashboard/{self.get_connector_name()}/redshift.yaml"))
        self.run_command(test_file_path=redshift_file_path)

    @staticmethod
    def get_connector_name() -> str:
        return "tableau"

    # ================================
    # FILTER CONFIGURATION METHODS
    # ================================

    def get_includes_dashboards(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.INCLUDE_DASHBOARDS

    def get_excludes_dashboards(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.EXCLUDE_DASHBOARDS

    def get_includes_charts(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.INCLUDE_CHARTS

    def get_excludes_charts(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.EXCLUDE_CHARTS

    def get_includes_datamodels(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.INCLUDE_DATAMODELS

    def get_excludes_datamodels(self) -> List[str]:  # noqa: UP006
        return TableauExpectedValues.EXCLUDE_DATAMODELS

    # ================================
    # EXPECTED COUNT METHODS
    # ================================

    def expected_dashboards_and_charts(self) -> int:
        return TableauExpectedValues.DASHBOARDS_AND_CHARTS

    def expected_lineage(self) -> int:
        return TableauExpectedValues.LINEAGE_EDGES

    def expected_tags(self) -> int:
        return TableauExpectedValues.TAGS

    def expected_datamodel_lineage(self) -> int:
        return TableauExpectedValues.DATAMODEL_LINEAGE

    def expected_datamodels(self) -> int:
        return TableauExpectedValues.DATAMODELS

    def expected_filtered_mix(self) -> int:
        return TableauExpectedValues.FILTERED_MIX

    def expected_filtered_sink_mix(self) -> int:
        return TableauExpectedValues.FILTERED_SINK_MIX

    def expected_dashboards_and_charts_after_patch(self) -> int:
        return TableauExpectedValues.DASHBOARDS_AND_CHARTS_AFTER_PATCH

    # ================================
    # ENHANCED TEST METHODS
    # ================================

    @pytest.mark.order(10)
    def test_metadata_ingestion_validation(self) -> None:
        """Test comprehensive metadata ingestion validation - runs after base ingestion"""
        # Validate dashboards
        self._validate_dashboard_metadata()

        # Validate charts
        self._validate_chart_metadata()

        # Validate data models
        self._validate_datamodel_metadata()

        # Validate tags
        self._validate_tags_metadata()

    @pytest.mark.order(12)
    def test_lineage_validation(self) -> None:
        """Test comprehensive lineage validation - Enable actual lineage testing"""
        # Remove the pytest.skip and implement actual lineage validation
        self._validate_dashboard_lineage()
        self._validate_datamodel_lineage_chain()

    @pytest.mark.order(13)
    def test_datamodel_content_validation(self) -> None:
        """Test data model content and structure validation"""
        self._validate_datamodel_fields()
        self._validate_datamodel_sql()
        self._validate_datamodel_types()

    @pytest.mark.order(14)
    def test_tags_and_metadata_quality(self) -> None:
        """Test tags and metadata quality"""
        self._validate_tag_assignment()
        self._validate_chart_types()
        self._validate_field_types()

    def _validate_dashboard_metadata(self) -> None:
        """Validate dashboard metadata completeness"""
        dashboards = self.openmetadata.list_entities(
            entity=Dashboard, params={"service": TableauExpectedValues.SERVICE_NAME}
        ).entities

        self.assertGreaterEqual(len(dashboards), len(TableauExpectedValues.EXPECTED_DASHBOARD_NAMES))

        dashboard_names = [dashboard.displayName for dashboard in dashboards]
        for expected_name in TableauExpectedValues.EXPECTED_DASHBOARD_NAMES:
            self.assertIn(
                expected_name,
                dashboard_names,
                f"Expected dashboard '{expected_name}' not found in ingested dashboards",
            )

        # Validate specific dashboard metadata
        analytics_dashboard = self.get_entity_by_name(Dashboard, "Analytics Workbook")
        if analytics_dashboard:
            self.assertIsNotNone(analytics_dashboard.fullyQualifiedName)
            self.assertIsNotNone(analytics_dashboard.service)
            # Validate that charts are properly linked
            if hasattr(analytics_dashboard, "charts") and analytics_dashboard.charts:
                self.assertGreaterEqual(
                    len(analytics_dashboard.charts.root),
                    4,
                    "Analytics Workbook should have 4 or more charts",
                )

    def _validate_chart_metadata(self) -> None:
        """Validate chart metadata completeness"""
        charts = self.openmetadata.list_entities(
            entity=Chart, params={"service": TableauExpectedValues.SERVICE_NAME}
        ).entities

        self.assertGreaterEqual(len(charts), len(TableauExpectedValues.EXPECTED_CHART_NAMES))

        chart_names = [chart.displayName for chart in charts]
        for expected_name in TableauExpectedValues.EXPECTED_CHART_NAMES:
            self.assertIn(
                expected_name,
                chart_names,
                f"Expected chart '{expected_name}' not found in ingested charts",
            )

    def _validate_datamodel_metadata(self) -> None:
        """Validate data model metadata completeness"""
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        # Should have at least one "Sales Summary" data model
        datamodel_names = [dm.displayName for dm in datamodels]
        self.assertIn("Sales Summary", datamodel_names, "Sales Summary data model not found")

        # Validate data model types
        datamodel_types = []
        for dm in datamodels:
            if hasattr(dm, "dataModelType") and dm.dataModelType:
                datamodel_types.append(str(dm.dataModelType))  # noqa: PERF401

        for expected_type in TableauExpectedValues.EXPECTED_DATAMODEL_TYPES:
            self.assertIn(
                expected_type,
                datamodel_types,
                f"Expected data model type '{expected_type}' not found",
            )

    def _validate_tags_metadata(self) -> None:
        """Validate tags metadata"""
        # Get all entities and check for expected tags
        all_entities = []
        all_entities.extend(
            self.openmetadata.list_entities(
                entity=Dashboard,
                params={"service": TableauExpectedValues.SERVICE_NAME},
                fields=["tags"],
            ).entities
        )
        all_entities.extend(
            self.openmetadata.list_entities(
                entity=Chart,
                params={"service": TableauExpectedValues.SERVICE_NAME},
                fields=["tags"],
            ).entities
        )

        found_tags = set()
        for entity in all_entities:
            if hasattr(entity, "tags") and entity.tags:
                for tag in entity.tags:
                    if hasattr(tag, "name"):
                        # Extract tag name from FQN
                        found_tags.add(str(tag.name))

        for expected_tag in TableauExpectedValues.EXPECTED_TAGS:
            self.assertIn(
                expected_tag,
                found_tags,
                f"Expected tag '{expected_tag}' not found in any entity",
            )

    def _validate_dashboard_lineage(self) -> None:
        """Validate dashboard lineage according to the knowledge base"""
        # Lineage chain: Tables -> TableauPublishedDatasource -> TableauEmbeddedDatasource -> Dashboard
        analytics_dashboard = self.get_entity_by_name(Dashboard, "Analytics Workbook")
        if analytics_dashboard:
            lineage = self.openmetadata.get_lineage_by_name(
                entity=Dashboard,
                fqn=analytics_dashboard.fullyQualifiedName.root,
                up_depth=5,  # Increased depth to capture full lineage chain
                down_depth=1,
            )

            if lineage and lineage.get("upstreamEdges"):
                self.assertGreater(
                    len(lineage["upstreamEdges"]),
                    6,
                    "Analytics Workbook should have upstream lineage",
                )

    def _validate_datamodel_lineage_chain(self) -> None:
        """Validate the complete lineage chain"""
        # Get Sales Summary data models
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        sales_summary_models = [dm for dm in datamodels if dm.displayName == "Sales Summary"]
        for datamodel in sales_summary_models:
            lineage = self.openmetadata.get_lineage_by_name(
                entity=DashboardDataModel,
                fqn=datamodel.fullyQualifiedName.root,
                up_depth=3,
                down_depth=3,
            )

            if lineage and lineage.get("upstreamEdges"):
                # Check for upstream connections (should connect to tables)
                for edge in lineage["upstreamEdges"]:
                    if lineage_query := edge["lineageDetails"].get("sqlQuery"):
                        self.assertEqual(
                            " ".join(lineage_query.split()),
                            " ".join(TableauExpectedValues.EXPECTED_DATAMODEL_SQL.split()),
                            "Lineage SQL query does't match expected SQL query",
                        )

    def _validate_datamodel_fields(self) -> None:
        """Validate data model fields/columns"""
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        sales_summary_models = [dm for dm in datamodels if dm.name.root == "Sales Summary"]

        for datamodel in sales_summary_models:
            if hasattr(datamodel, "columns") and datamodel.columns:
                column_names = [col.name.root for col in datamodel.columns]

                # Check for expected fields
                found_fields = 0
                for expected_field in TableauExpectedValues.EXPECTED_DATAMODEL_FIELDS:
                    if expected_field in column_names:
                        found_fields += 1

                self.assertGreater(
                    found_fields,
                    5,  # Should find at least 5 expected fields
                    f"Data model {datamodel.name.root} should contain expected fields",
                )

    def _validate_datamodel_sql(self) -> None:
        """Validate data model SQL content"""
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        sales_summary_models = [dm for dm in datamodels if dm.name.root == "Sales Summary"]

        for datamodel in sales_summary_models:
            if hasattr(datamodel, "sql") and datamodel.sql:
                sql_content = datamodel.sql.root if hasattr(datamodel.sql, "root") else str(datamodel.sql)

                # Check for key SQL elements
                self.assertIn("SELECT", sql_content.upper(), "SQL should contain SELECT statement")
                self.assertIn("JOIN", sql_content.upper(), "SQL should contain JOIN statements")

                # Check for expected table references
                for table in TableauExpectedValues.EXPECTED_SOURCE_TABLES:
                    self.assertIn(
                        table,
                        sql_content.lower(),
                        f"SQL should reference table '{table}'",
                    )

    def _validate_datamodel_types(self) -> None:
        """Validate data model types"""
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        found_types = set()
        for datamodel in datamodels:
            if hasattr(datamodel, "dataModelType") and datamodel.dataModelType:
                found_types.add(str(datamodel.dataModelType))

        for expected_type in TableauExpectedValues.EXPECTED_DATAMODEL_TYPES:
            self.assertIn(
                expected_type,
                found_types,
                f"Expected data model type '{expected_type}' not found",
            )

    def _validate_tag_assignment(self) -> None:
        """Validate tag assignment to entities"""
        # Check Analytics Workbook dashboard
        analytics_dashboard = self.get_entity_by_name(Dashboard, "Analytics Workbook")
        if analytics_dashboard and hasattr(analytics_dashboard, "tags") and analytics_dashboard.tags:
            dashboard_tags = {str(tag.name) for tag in analytics_dashboard.tags}
            for expected_tag in TableauExpectedValues.EXPECTED_TAGS:
                self.assertIn(
                    expected_tag,
                    dashboard_tags,
                    f"Dashboard should have tag '{expected_tag}'",
                )

        # Check charts
        for chart_name in TableauExpectedValues.EXPECTED_CHART_NAMES:
            chart = self.get_entity_by_name(Chart, chart_name, fields=None)
            if chart and hasattr(chart, "tags") and chart.tags:
                chart_tags = {str(tag.name) for tag in chart.tags}
                for expected_tag in TableauExpectedValues.EXPECTED_TAGS:
                    self.assertIn(
                        expected_tag,
                        chart_tags,
                        f"Chart '{chart_name}' should have tag '{expected_tag}'",
                    )

    def _validate_chart_types(self) -> None:
        """Validate chart types"""
        charts = self.openmetadata.list_entities(
            entity=Chart, params={"service": TableauExpectedValues.SERVICE_NAME}
        ).entities

        for chart in charts:
            if hasattr(chart, "chartType") and chart.chartType:
                self.assertEqual(
                    str(chart.chartType),
                    TableauExpectedValues.EXPECTED_CHART_TYPE,
                    f"Chart '{chart.name.root}' should have type '{TableauExpectedValues.EXPECTED_CHART_TYPE}'",
                )

    def _validate_field_types(self) -> None:
        """Validate field types in data models"""
        datamodels = self.openmetadata.list_entities(
            entity=DashboardDataModel,
            params={"service": TableauExpectedValues.SERVICE_NAME},
        ).entities

        for datamodel in datamodels:
            if hasattr(datamodel, "columns") and datamodel.columns:
                for column in datamodel.columns:
                    if hasattr(column, "dataType") and column.dataType:
                        # Check if field type matches expected
                        field_type = str(column.dataType)
                        self.assertIn(
                            "DataType.RECORD",
                            field_type,
                            f"Field '{column.name.root}' should have tableau-related type",
                        )

    @pytest.mark.order(11)
    def test_lineage(self) -> None:
        """Enable lineage testing - Remove the skip"""
        self._validate_dashboard_lineage()
        self._validate_datamodel_lineage_chain()

    def assert_not_including(self, source_status: Status, sink_status: Status):
        """
        Override base method for Tableau-specific behavior.
        """
        self.assertTrue(len(source_status.failures) == 0)
        self.assertTrue(len(source_status.warnings) == 0)
        self.assertTrue(len(source_status.filtered) >= 5)

        # We can have a diff of 1 element if we are counting the service, which is only marked as ingested in the
        # first go
        self.assertTrue(
            self.expected_dashboards_and_charts() <= (len(source_status.records) + len(source_status.updated_records))
        )
        self.assertTrue(len(sink_status.failures) == 0)
        self.assertTrue(len(sink_status.warnings) == 0)
        self.assertTrue(
            self.expected_dashboards_and_charts() <= (len(sink_status.records) + len(sink_status.updated_records))
        )

    def assert_for_vanilla_ingestion(self, source_status: Status, sink_status: Status) -> None:
        self.assertTrue(len(source_status.failures) == 0)
        self.assertTrue(len(source_status.warnings) == 0)
        self.assertTrue(len(source_status.filtered) >= 5)
        self.assertGreaterEqual(
            (len(source_status.records) + len(source_status.updated_records)),
            self.expected_dashboards_and_charts_after_patch()
            + self.expected_tags()
            + self.expected_lineage()
            + self.expected_datamodel_lineage(),
        )
        self.assertTrue(len(sink_status.failures) == 0)
        self.assertTrue(len(sink_status.warnings) == 0)
        self.assertGreaterEqual(
            (len(sink_status.records) + len(sink_status.updated_records)),
            self.expected_dashboards_and_charts_after_patch() + self.expected_tags() + self.expected_datamodels(),
        )

    def get_entity_by_name(
        self,
        entity_type,
        name: str,
        service: str = TableauExpectedValues.SERVICE_NAME,
        fields: List = ["tags", "charts"],  # noqa: B006, UP006
    ):
        """Helper to get entity by name or displayName"""
        entities = self.openmetadata.list_entities(
            entity=entity_type, params={"service": service}, fields=fields
        ).entities

        for entity in entities:
            # Check both name and displayName for matches
            entity_name = entity.name.root if hasattr(entity.name, "root") else str(entity.name)
            entity_display_name = (
                entity.displayName.root
                if hasattr(entity, "displayName") and entity.displayName and hasattr(entity.displayName, "root")
                else (str(entity.displayName) if hasattr(entity, "displayName") and entity.displayName else None)
            )

            if entity_name == name or entity_display_name == name:  # noqa: PLR1714
                return entity
        return None

    def validate_entity_exists(self, entity_type, name: str, service: str = TableauExpectedValues.SERVICE_NAME):
        """Helper to validate entity exists"""
        entity = self.get_entity_by_name(entity_type, name, service)
        self.assertIsNotNone(entity, f"{entity_type.__name__} '{name}' not found in service '{service}'")
        return entity