OpenMetadata/ingestion/tests/unit/sdk/test_custom_properties.py
Sriharsha Chintalapani bb1395fc72
Implement Modern Fluent API Pattern for OpenMetadata Java Client (#23239)
* Implement Modern Fluent API Pattern for OpenMetadata Java Client

* Add Lineage, Bulk, Search static methods

* Add all API support for Java & Python SDKs

* Add Python SDKs and mock tests

* Add Fluent APIs for sdks

* Add Fluent APIs for sdks

* Add Fluent APIs for sdks, support async import/export

* Remove unnecessary scripts

* fix py checkstyle

* fix tests with new plural form sdks

* Fix tests

* remove examples from python sdk

* remove examples from python sdk

* Fix type check

* Fix pyformat check

* Fix pyformat check

* fix python integration tests

* fix pycheck and pytests

* fix search api pycheck

* fix pycheck

* fix pycheck

* fix pycheck

* Fix test_sdk_integration

* Improvements to SDK

* Remove SDK coverage for Python 3.9

* Remove SDK coverage for Python 3.9

* Remove SDK coverage for Python 3.9
2025-09-29 16:07:02 -07:00

322 lines
12 KiB
Python

"""
Unit tests for custom property operations in SDK.
"""
from unittest.mock import Mock, patch
from uuid import UUID
from metadata.generated.schema.entity.data.glossary import Glossary
from metadata.generated.schema.entity.data.table import Table
from metadata.generated.schema.type import basic
from metadata.sdk.entities.custom_properties import (
CustomProperties,
CustomPropertyUpdater,
GlossaryCustomProperties,
TableCustomProperties,
)
class TestCustomPropertyUpdater:
"""Test CustomPropertyUpdater class."""
def test_with_property(self):
"""Test adding a single property."""
updater = CustomPropertyUpdater(Table, "test-id")
updater.with_property("key1", "value1")
assert updater.properties == {"key1": "value1"}
def test_with_properties(self):
"""Test adding multiple properties."""
updater = CustomPropertyUpdater(Table, "test-id")
properties = {"key1": "value1", "key2": "value2"}
updater.with_properties(properties)
assert updater.properties == properties
def test_clear_property(self):
"""Test clearing a specific property."""
updater = CustomPropertyUpdater(Table, "test-id")
updater.with_property("key1", "value1")
updater.clear_property("key1")
assert updater.properties == {"key1": None}
def test_clear_all(self):
"""Test clearing all properties."""
updater = CustomPropertyUpdater(Table, "test-id")
updater.with_property("key1", "value1")
updater.clear_all()
assert updater.clear_all_flag is True
@patch("metadata.sdk.entities.custom_properties.CustomPropertyUpdater._get_client")
def test_execute_with_new_properties(self, mock_get_client):
"""Test executing update with new properties."""
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity
mock_table = Mock(spec=Table)
mock_table.extension = None
mock_table.model_copy = Mock(return_value=mock_table)
mock_client.get_by_id.return_value = mock_table
# Create updated entity
updated_table = Mock(spec=Table)
mock_client.patch.return_value = updated_table
# Execute update
updater = CustomPropertyUpdater(Table, "test-id")
updater.with_property("businessImportance", "HIGH")
result = updater.execute()
# Verify
mock_client.get_by_id.assert_called_once_with(
entity=Table, entity_id="test-id", fields=["extension"]
)
mock_client.patch.assert_called_once()
assert result == updated_table
@patch("metadata.sdk.entities.custom_properties.CustomPropertyUpdater._get_client")
def test_execute_with_existing_extension(self, mock_get_client):
"""Test updating entity with existing extension."""
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity with existing extension
mock_table = Mock(spec=Table)
existing_extension = basic.EntityExtension(
root={"existingKey": "existingValue"}
)
mock_table.extension = existing_extension
mock_table_copy = Mock(spec=Table)
mock_table.model_copy = Mock(return_value=mock_table_copy)
mock_client.get_by_id.return_value = mock_table
# Create updated entity
updated_table = Mock(spec=Table)
mock_client.patch.return_value = updated_table
# Execute update
updater = CustomPropertyUpdater(Table, "test-id")
updater.with_property("newKey", "newValue")
result = updater.execute()
# Verify the extension was updated correctly
mock_client.patch.assert_called_once()
args = mock_client.patch.call_args
assert args[1]["destination"].extension.root == {
"existingKey": "existingValue",
"newKey": "newValue",
}
@patch("metadata.sdk.entities.custom_properties.CustomPropertyUpdater._get_client")
def test_execute_clear_all(self, mock_get_client):
"""Test clearing all custom properties."""
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity with existing extension
mock_table = Mock(spec=Table)
existing_extension = basic.EntityExtension(
root={"key1": "value1", "key2": "value2"}
)
mock_table.extension = existing_extension
mock_table_copy = Mock(spec=Table)
mock_table_copy.extension = None
mock_table.model_copy = Mock(return_value=mock_table_copy)
mock_client.get_by_id.return_value = mock_table
# Create updated entity
updated_table = Mock(spec=Table)
mock_client.patch.return_value = updated_table
# Execute clear all
updater = CustomPropertyUpdater(Table, "test-id")
updater.clear_all()
result = updater.execute()
# Verify extension was cleared
mock_client.patch.assert_called_once()
args = mock_client.patch.call_args
assert args[1]["destination"].extension is None
@patch("metadata.sdk.entities.custom_properties.CustomPropertyUpdater._get_client")
def test_execute_by_fqn(self, mock_get_client):
"""Test updating entity by FQN."""
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity
mock_table = Mock(spec=Table)
mock_table.extension = None
mock_table.model_copy = Mock(return_value=mock_table)
mock_client.get_by_name.return_value = mock_table
# Create updated entity
updated_table = Mock(spec=Table)
mock_client.patch.return_value = updated_table
# Execute update by FQN
updater = CustomPropertyUpdater(
Table, "service.database.schema.table", is_fqn=True
)
updater.with_property("key", "value")
result = updater.execute()
# Verify get_by_name was called
mock_client.get_by_name.assert_called_once_with(
entity=Table, fqn="service.database.schema.table", fields=["extension"]
)
class TestCustomProperties:
"""Test CustomProperties factory class."""
def test_update(self):
"""Test creating updater by ID."""
updater = CustomProperties.update(Table, "test-id")
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Table
assert updater.identifier == "test-id"
assert updater.is_fqn is False
def test_update_with_uuid(self):
"""Test creating updater with UUID."""
test_uuid = UUID("12345678-1234-5678-1234-567812345678")
updater = CustomProperties.update(Table, test_uuid)
assert updater.identifier == str(test_uuid)
def test_update_by_name(self):
"""Test creating updater by name/FQN."""
updater = CustomProperties.update_by_name(
Table, "service.database.schema.table"
)
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Table
assert updater.identifier == "service.database.schema.table"
assert updater.is_fqn is True
class TestTableCustomProperties:
"""Test Table-specific custom property methods."""
def test_update(self):
"""Test creating table updater by ID."""
updater = TableCustomProperties.update("test-id")
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Table
assert updater.identifier == "test-id"
def test_update_by_name(self):
"""Test creating table updater by FQN."""
updater = TableCustomProperties.update_by_name("service.database.schema.table")
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Table
assert updater.identifier == "service.database.schema.table"
assert updater.is_fqn is True
class TestGlossaryCustomProperties:
"""Test Glossary-specific custom property methods."""
def test_update(self):
"""Test creating glossary updater by ID."""
updater = GlossaryCustomProperties.update("test-id")
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Glossary
assert updater.identifier == "test-id"
def test_update_by_name(self):
"""Test creating glossary updater by name."""
updater = GlossaryCustomProperties.update_by_name("BusinessGlossary")
assert isinstance(updater, CustomPropertyUpdater)
assert updater.entity_type == Glossary
assert updater.identifier == "BusinessGlossary"
assert updater.is_fqn is True
class TestIntegrationWithBaseEntity:
"""Test integration with BaseEntity class."""
@patch("metadata.sdk.entities.base.BaseEntity._get_client")
def test_update_custom_properties_from_base_entity(self, mock_get_client):
"""Test using custom properties through BaseEntity."""
from metadata.sdk import Tables
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity
mock_table = Mock(spec=Table)
mock_table.extension = None
mock_table.model_copy = Mock(return_value=mock_table)
mock_client.get_by_id.return_value = mock_table
mock_client.patch.return_value = mock_table
# Use BaseEntity method
updater = Tables.update_custom_properties("test-id")
assert isinstance(updater, CustomPropertyUpdater)
# Execute update
result = updater.with_property("key", "value").execute()
assert result == mock_table
@patch("metadata.sdk.entities.base.BaseEntity._get_client")
def test_update_custom_properties_by_name_from_base_entity(self, mock_get_client):
"""Test using custom properties by name through BaseEntity."""
from metadata.sdk import Tables
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity
mock_table = Mock(spec=Table)
mock_table.extension = None
mock_table.model_copy = Mock(return_value=mock_table)
mock_client.get_by_name.return_value = mock_table
mock_client.patch.return_value = mock_table
# Use BaseEntity method
updater = Tables.update_custom_properties_by_name(
"service.database.schema.table"
)
assert isinstance(updater, CustomPropertyUpdater)
# Execute update
result = updater.with_property("key", "value").execute()
assert result == mock_table
class TestFluentChaining:
"""Test fluent API chaining."""
@patch("metadata.sdk.entities.custom_properties.CustomPropertyUpdater._get_client")
def test_fluent_chaining(self, mock_get_client):
"""Test chaining multiple operations."""
# Setup mock client
mock_client = Mock()
mock_get_client.return_value = mock_client
# Setup mock entity
mock_table = Mock(spec=Table)
mock_table.extension = None
mock_table.model_copy = Mock(return_value=mock_table)
mock_client.get_by_id.return_value = mock_table
mock_client.patch.return_value = mock_table
# Test fluent chaining
result = (
CustomProperties.update(Table, "test-id")
.with_property("key1", "value1")
.with_property("key2", "value2")
.with_properties({"key3": "value3", "key4": "value4"})
.clear_property("key1")
.execute()
)
assert result == mock_table
mock_client.patch.assert_called_once()