OpenMetadata/ingestion/tests/unit/sdk/test_sdk_fluent_api.py

"""
Integration tests for SDK fluent API with plural entity classes
"""
from unittest.mock import MagicMock, patch
from uuid import UUID

from metadata.generated.schema.api.data.createGlossary import CreateGlossaryRequest
from metadata.generated.schema.api.data.createGlossaryTerm import (
    CreateGlossaryTermRequest,
)
from metadata.generated.schema.api.data.createTable import CreateTableRequest
from metadata.generated.schema.api.teams.createTeam import CreateTeamRequest
from metadata.generated.schema.api.teams.createUser import CreateUserRequest
from metadata.generated.schema.entity.data.database import Database
from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
from metadata.generated.schema.entity.data.glossary import Glossary
from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
from metadata.generated.schema.entity.data.table import Column, DataType, Table
from metadata.generated.schema.entity.teams.team import Team
from metadata.generated.schema.entity.teams.user import User

# Import SDK plural entity classes
from metadata.sdk import (
    Databases,
    DatabaseSchemas,
    Glossaries,
    GlossaryTerms,
    Tables,
    Teams,
    Users,
)


class TestSDKFluentAPI:
    """Test SDK fluent API with plural entity classes"""

    def setup_method(self):
        """Set up test fixtures"""
        self.mock_ometa = MagicMock()

        # Mock the client getter to return our mock
        with patch.object(Tables, "_get_client", return_value=self.mock_ometa):
            with patch.object(Users, "_get_client", return_value=self.mock_ometa):
                with patch.object(
                    Databases, "_get_client", return_value=self.mock_ometa
                ):
                    with patch.object(
                        DatabaseSchemas, "_get_client", return_value=self.mock_ometa
                    ):
                        with patch.object(
                            Teams, "_get_client", return_value=self.mock_ometa
                        ):
                            with patch.object(
                                Glossaries, "_get_client", return_value=self.mock_ometa
                            ):
                                with patch.object(
                                    GlossaryTerms,
                                    "_get_client",
                                    return_value=self.mock_ometa,
                                ):
                                    self.setup_mocks()

    def setup_mocks(self):
        """Setup mock responses"""
        # Mock IDs
        self.table_id = UUID("550e8400-e29b-41d4-a716-446655440000")
        self.user_id = UUID("650e8400-e29b-41d4-a716-446655440000")
        self.database_id = UUID("750e8400-e29b-41d4-a716-446655440000")
        self.schema_id = UUID("850e8400-e29b-41d4-a716-446655440000")
        self.team_id = UUID("950e8400-e29b-41d4-a716-446655440000")
        self.glossary_id = UUID("a50e8400-e29b-41d4-a716-446655440000")
        self.term_id = UUID("b50e8400-e29b-41d4-a716-446655440000")

    @patch.object(Tables, "_get_client")
    def test_tables_create(self, mock_get_client):
        """Test creating a table using Tables SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange
        columns = [
            Column(name="id", dataType=DataType.INT),
            Column(name="name", dataType=DataType.VARCHAR, dataLength=255),
        ]

        create_request = CreateTableRequest(
            name="test_table",
            databaseSchema="test_schema",
            columns=columns,
        )

        mock_table = MagicMock(spec=Table)
        mock_table.id = self.table_id
        mock_table.name = "test_table"
        mock_table.fullyQualifiedName = "service.database.schema.test_table"
        self.mock_ometa.create_or_update.return_value = mock_table

        # Act
        result = Tables.create(create_request)

        # Assert
        assert result.name == "test_table"
        self.mock_ometa.create_or_update.assert_called_once_with(create_request)

    @patch.object(Tables, "_get_client")
    def test_tables_retrieve_by_name(self, mock_get_client):
        """Test retrieving a table by name using Tables SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange
        mock_table = MagicMock(spec=Table)
        mock_table.id = self.table_id
        mock_table.name = "test_table"
        mock_table.fullyQualifiedName = "service.database.schema.test_table"
        self.mock_ometa.get_by_name.return_value = mock_table

        # Act
        result = Tables.retrieve_by_name("service.database.schema.test_table")

        # Assert
        assert result.name == "test_table"
        self.mock_ometa.get_by_name.assert_called_once_with(
            entity=Table, fqn="service.database.schema.test_table", fields=None
        )

    @patch.object(Users, "_get_client")
    def test_users_create(self, mock_get_client):
        """Test creating a user using Users SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange
        create_request = CreateUserRequest(
            name="test_user",
            email="test@example.com",
        )

        mock_user = MagicMock(spec=User)
        mock_user.id = self.user_id
        mock_user.name = "test_user"
        mock_user.email = "test@example.com"
        self.mock_ometa.create_or_update.return_value = mock_user

        # Act
        result = Users.create(create_request)

        # Assert
        assert result.name == "test_user"
        assert result.email == "test@example.com"
        self.mock_ometa.create_or_update.assert_called_once_with(create_request)

    @patch.object(Databases, "_get_client")
    def test_databases_list(self, mock_get_client):
        """Test listing databases using Databases SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange
        mock_db1 = MagicMock(spec=Database)
        mock_db1.name = "db1"
        mock_db2 = MagicMock(spec=Database)
        mock_db2.name = "db2"

        mock_response = MagicMock()
        mock_response.entities = [mock_db1, mock_db2]
        self.mock_ometa.list_entities.return_value = mock_response

        # Act
        result = Databases.list()

        # Assert
        assert len(result.entities) == 2
        assert result.entities[0].name == "db1"
        assert result.entities[1].name == "db2"
        self.mock_ometa.list_entities.assert_called_once_with(
            entity=Database, fields=None, after=None, before=None, limit=10, params=None
        )

    @patch.object(Teams, "_get_client")
    def test_teams_create_and_delete(self, mock_get_client):
        """Test creating and deleting a team using Teams SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange - Create
        create_request = CreateTeamRequest(
            name="test_team",
            teamType="Group",
        )

        mock_team = MagicMock(spec=Team)
        mock_team.id = self.team_id
        mock_team.name = "test_team"
        self.mock_ometa.create_or_update.return_value = mock_team

        # Act - Create
        created_team = Teams.create(create_request)

        # Assert - Create
        assert created_team.name == "test_team"

        # Arrange - Delete
        self.mock_ometa.delete.return_value = None

        # Act - Delete
        Teams.delete(str(self.team_id))

        # Assert - Delete
        self.mock_ometa.delete.assert_called_once_with(
            entity=Team, entity_id=str(self.team_id), recursive=False, hard_delete=False
        )

    @patch.object(Glossaries, "_get_client")
    @patch.object(GlossaryTerms, "_get_client")
    def test_glossary_and_terms_workflow(self, mock_terms_client, mock_glossary_client):
        """Test creating glossary and terms using SDK"""
        mock_glossary_client.return_value = self.mock_ometa
        mock_terms_client.return_value = self.mock_ometa

        # Create Glossary
        glossary_request = CreateGlossaryRequest(
            name="business_glossary",
            displayName="Business Glossary",
            description="Company business terms",
        )

        mock_glossary = MagicMock(spec=Glossary)
        mock_glossary.id = self.glossary_id
        mock_glossary.name = "business_glossary"
        self.mock_ometa.create_or_update.return_value = mock_glossary

        created_glossary = Glossaries.create(glossary_request)
        assert created_glossary.name == "business_glossary"

        # Create Glossary Term
        term_request = CreateGlossaryTermRequest(
            name="customer",
            displayName="Customer",
            description="A person or organization that buys goods or services",
            glossary="business_glossary",  # Use FQN string
        )

        mock_term = MagicMock(spec=GlossaryTerm)
        mock_term.id = self.term_id
        mock_term.name = "customer"
        mock_term.glossary = mock_glossary
        self.mock_ometa.create_or_update.return_value = mock_term

        created_term = GlossaryTerms.create(term_request)
        assert created_term.name == "customer"

    @patch.object(DatabaseSchemas, "_get_client")
    def test_database_schemas_search(self, mock_get_client):
        """Test searching database schemas using DatabaseSchemas SDK"""
        mock_get_client.return_value = self.mock_ometa

        # Arrange
        mock_schema1 = MagicMock(spec=DatabaseSchema)
        mock_schema1.name = "public"
        mock_schema2 = MagicMock(spec=DatabaseSchema)
        mock_schema2.name = "staging"

        self.mock_ometa.es_search_from_fqn.return_value = [mock_schema1, mock_schema2]

        # Act
        results = DatabaseSchemas.search("staging")

        # Assert
        assert len(results) == 2
        self.mock_ometa.es_search_from_fqn.assert_called_once_with(
            entity_type=DatabaseSchema, fqn_search_string="staging", size=10
        )


class TestSDKEntityTypes:
    """Test that all SDK entities have proper entity_type methods"""

    def test_all_entities_have_entity_type(self):
        """Verify all plural SDK entities implement entity_type method"""
        from metadata.sdk import (
            APICollections,
            APIEndpoints,
            Charts,
            Classifications,
            Containers,
            DashboardDataModels,
            Dashboards,
            Databases,
            DatabaseSchemas,
            DataContracts,
            DataProducts,
            Domains,
            Glossaries,
            GlossaryTerms,
            Metrics,
            MLModels,
            Pipelines,
            Queries,
            SearchIndexes,
            StoredProcedures,
            Tables,
            Tags,
            Teams,
            TestCases,
            TestDefinitions,
            TestSuites,
            Users,
        )

        entities = [
            Tables,
            Users,
            Databases,
            DatabaseSchemas,
            Dashboards,
            Charts,
            Pipelines,
            MLModels,
            Containers,
            Queries,
            Metrics,
            Glossaries,
            GlossaryTerms,
            Classifications,
            Tags,
            Domains,
            DataProducts,
            DataContracts,
            APICollections,
            APIEndpoints,
            SearchIndexes,
            StoredProcedures,
            DashboardDataModels,
            TestCases,
            TestDefinitions,
            TestSuites,
            Teams,
        ]

        for entity_class in entities:
            # Check that entity_type method exists and returns a type
            assert hasattr(
                entity_class, "entity_type"
            ), f"{entity_class.__name__} missing entity_type method"
            entity_type = entity_class.entity_type()
            assert (
                entity_type is not None
            ), f"{entity_class.__name__}.entity_type() returned None"
            assert isinstance(
                entity_type, type
            ), f"{entity_class.__name__}.entity_type() did not return a type"