mirror of
				https://github.com/open-metadata/OpenMetadata.git
				synced 2025-11-04 12:36:23 +00:00 
			
		
		
		
	* 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
		
			
				
	
	
		
			294 lines
		
	
	
		
			10 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			294 lines
		
	
	
		
			10 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
"""
 | 
						|
Comprehensive unit tests for Container entity with full mock coverage.
 | 
						|
"""
 | 
						|
import unittest
 | 
						|
from unittest.mock import MagicMock
 | 
						|
from uuid import UUID
 | 
						|
 | 
						|
from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 | 
						|
from metadata.generated.schema.entity.data.container import Container as ContainerEntity
 | 
						|
from metadata.generated.schema.entity.data.container import (
 | 
						|
    ContainerDataModel,
 | 
						|
    FileFormat,
 | 
						|
)
 | 
						|
from metadata.generated.schema.type.entityReference import EntityReference
 | 
						|
from metadata.sdk import Containers
 | 
						|
 | 
						|
 | 
						|
class TestContainerEntity(unittest.TestCase):
 | 
						|
    """Comprehensive tests for Container entity operations"""
 | 
						|
 | 
						|
    def setUp(self):
 | 
						|
        """Set up test fixtures"""
 | 
						|
        self.mock_ometa = MagicMock()
 | 
						|
 | 
						|
        # Set default client directly
 | 
						|
        Containers.set_default_client(self.mock_ometa)
 | 
						|
 | 
						|
        # Test data
 | 
						|
        self.container_id = "650e8400-e29b-41d4-a716-446655440000"
 | 
						|
        self.container_fqn = "s3-prod.analytics-bucket"
 | 
						|
 | 
						|
    def test_create_container(self):
 | 
						|
        """Test creating a container"""
 | 
						|
        # Arrange
 | 
						|
        create_request = CreateContainerRequest(
 | 
						|
            name="analytics-bucket",
 | 
						|
            service="s3-prod",
 | 
						|
            displayName="Analytics Bucket",
 | 
						|
            description="S3 bucket for analytics data",
 | 
						|
            fileFormats=[FileFormat.parquet, FileFormat.csv],
 | 
						|
        )
 | 
						|
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.name = "analytics-bucket"
 | 
						|
        expected_container.fullyQualifiedName = self.container_fqn
 | 
						|
        expected_container.fileFormats = [FileFormat.parquet, FileFormat.csv]
 | 
						|
 | 
						|
        self.mock_ometa.create_or_update.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.create(create_request)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(str(result.id), self.container_id)
 | 
						|
        self.assertEqual(result.name, "analytics-bucket")
 | 
						|
        self.assertEqual(result.fullyQualifiedName, self.container_fqn)
 | 
						|
        self.assertEqual(len(result.fileFormats), 2)
 | 
						|
        self.mock_ometa.create_or_update.assert_called_once_with(create_request)
 | 
						|
 | 
						|
    def test_retrieve_container_by_id(self):
 | 
						|
        """Test retrieving a container by ID"""
 | 
						|
        # Arrange
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.name = "analytics-bucket"
 | 
						|
        expected_container.description = "Analytics data storage"
 | 
						|
 | 
						|
        self.mock_ometa.get_by_id.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve(self.container_id)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(str(result.id), self.container_id)
 | 
						|
        self.assertEqual(result.name, "analytics-bucket")
 | 
						|
        self.mock_ometa.get_by_id.assert_called_once_with(
 | 
						|
            entity=ContainerEntity, entity_id=self.container_id, fields=None
 | 
						|
        )
 | 
						|
 | 
						|
    def test_retrieve_container_with_children(self):
 | 
						|
        """Test retrieving container with child containers"""
 | 
						|
        # Arrange
 | 
						|
        fields = ["children", "parent", "dataModel"]
 | 
						|
 | 
						|
        # Mock child containers
 | 
						|
        child1 = EntityReference(
 | 
						|
            id=UUID("750e8400-e29b-41d4-a716-446655440001"),
 | 
						|
            type="container",
 | 
						|
            name="raw-data",
 | 
						|
        )
 | 
						|
        child2 = EntityReference(
 | 
						|
            id=UUID("750e8400-e29b-41d4-a716-446655440002"),
 | 
						|
            type="container",
 | 
						|
            name="processed-data",
 | 
						|
        )
 | 
						|
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.name = "analytics-bucket"
 | 
						|
        expected_container.children = [child1, child2]
 | 
						|
 | 
						|
        self.mock_ometa.get_by_id.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve(self.container_id, fields=fields)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertIsNotNone(result.children)
 | 
						|
        self.assertEqual(len(result.children), 2)
 | 
						|
        self.assertEqual(result.children[0].name, "raw-data")
 | 
						|
        self.mock_ometa.get_by_id.assert_called_once_with(
 | 
						|
            entity=ContainerEntity, entity_id=self.container_id, fields=fields
 | 
						|
        )
 | 
						|
 | 
						|
    def test_retrieve_container_by_name(self):
 | 
						|
        """Test retrieving a container by name"""
 | 
						|
        # Arrange
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.name = "analytics-bucket"
 | 
						|
        expected_container.fullyQualifiedName = self.container_fqn
 | 
						|
 | 
						|
        self.mock_ometa.get_by_name.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve_by_name(self.container_fqn)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(result.fullyQualifiedName, self.container_fqn)
 | 
						|
        self.mock_ometa.get_by_name.assert_called_once_with(
 | 
						|
            entity=ContainerEntity, fqn=self.container_fqn, fields=None
 | 
						|
        )
 | 
						|
 | 
						|
    def test_update_container(self):
 | 
						|
        """Test updating a container"""
 | 
						|
        # Arrange
 | 
						|
        container_to_update = MagicMock(spec=ContainerEntity)
 | 
						|
        container_to_update.id = UUID(self.container_id)
 | 
						|
        container_to_update.name = "analytics-bucket"
 | 
						|
        container_to_update.description = "Updated analytics bucket"
 | 
						|
 | 
						|
        # Mock the get_by_id to return the current state
 | 
						|
        current_entity = MagicMock(spec=type(container_to_update))
 | 
						|
        current_entity.id = (
 | 
						|
            container_to_update.id
 | 
						|
            if hasattr(container_to_update, "id")
 | 
						|
            else UUID(self.entity_id)
 | 
						|
        )
 | 
						|
        self.mock_ometa.get_by_id.return_value = current_entity
 | 
						|
 | 
						|
        # Mock the patch to return the updated entity
 | 
						|
        self.mock_ometa.patch.return_value = container_to_update
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.update(container_to_update)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(result.description, "Updated analytics bucket")
 | 
						|
        # Verify get_by_id was called to fetch current state
 | 
						|
        self.mock_ometa.get_by_id.assert_called_once()
 | 
						|
        # Verify patch was called with source and destination
 | 
						|
        self.mock_ometa.patch.assert_called_once()
 | 
						|
 | 
						|
    def test_delete_container(self):
 | 
						|
        """Test deleting a container"""
 | 
						|
        # Act
 | 
						|
        Containers.delete(self.container_id, recursive=True, hard_delete=False)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.mock_ometa.delete.assert_called_once_with(
 | 
						|
            entity=ContainerEntity,
 | 
						|
            entity_id=self.container_id,
 | 
						|
            recursive=True,
 | 
						|
            hard_delete=False,
 | 
						|
        )
 | 
						|
 | 
						|
    def test_container_with_data_model(self):
 | 
						|
        """Test container with data model"""
 | 
						|
        # Arrange
 | 
						|
        data_model = MagicMock(spec=ContainerDataModel)
 | 
						|
        data_model.isPartitioned = True
 | 
						|
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.dataModel = data_model
 | 
						|
 | 
						|
        self.mock_ometa.get_by_id.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve(self.container_id, fields=["dataModel"])
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertIsNotNone(result.dataModel)
 | 
						|
        self.assertEqual(result.dataModel.isPartitioned, True)
 | 
						|
 | 
						|
    def test_container_hierarchy(self):
 | 
						|
        """Test container with parent reference"""
 | 
						|
        # Arrange
 | 
						|
        parent_ref = EntityReference(
 | 
						|
            id=UUID("550e8400-e29b-41d4-a716-446655440000"),
 | 
						|
            type="container",
 | 
						|
            name="data-lake",
 | 
						|
        )
 | 
						|
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.parent = parent_ref
 | 
						|
 | 
						|
        self.mock_ometa.get_by_id.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve(self.container_id, fields=["parent"])
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertIsNotNone(result.parent)
 | 
						|
        self.assertEqual(result.parent.name, "data-lake")
 | 
						|
 | 
						|
    def test_container_with_prefix(self):
 | 
						|
        """Test container with prefix path"""
 | 
						|
        # Arrange
 | 
						|
        create_request = CreateContainerRequest(
 | 
						|
            name="partitioned-data",
 | 
						|
            service="s3-prod",
 | 
						|
            prefix="/data/year=2024/month=01/",
 | 
						|
        )
 | 
						|
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.name = "partitioned-data"
 | 
						|
        expected_container.prefix = "/data/year=2024/month=01/"
 | 
						|
 | 
						|
        self.mock_ometa.create_or_update.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.create(create_request)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(result.prefix, "/data/year=2024/month=01/")
 | 
						|
 | 
						|
    def test_list_containers(self):
 | 
						|
        """Test listing containers with pagination"""
 | 
						|
        # Arrange
 | 
						|
        mock_bucket1 = MagicMock(spec=ContainerEntity)
 | 
						|
        mock_bucket1.name = "bucket1"
 | 
						|
        mock_bucket2 = MagicMock(spec=ContainerEntity)
 | 
						|
        mock_bucket2.name = "bucket2"
 | 
						|
        mock_bucket3 = MagicMock(spec=ContainerEntity)
 | 
						|
        mock_bucket3.name = "bucket3"
 | 
						|
 | 
						|
        mock_response = MagicMock()
 | 
						|
        mock_response.entities = [mock_bucket1, mock_bucket2, mock_bucket3]
 | 
						|
 | 
						|
        self.mock_ometa.list_entities.return_value = mock_response
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.list(limit=25, after="cursor456")
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(len(result.entities), 3)
 | 
						|
        self.assertEqual(result.entities[0].name, "bucket1")
 | 
						|
        self.mock_ometa.list_entities.assert_called_once()
 | 
						|
 | 
						|
    def test_container_size_and_objects(self):
 | 
						|
        """Test container with size and object count"""
 | 
						|
        # Arrange
 | 
						|
        expected_container = MagicMock(spec=ContainerEntity)
 | 
						|
        expected_container.id = UUID(self.container_id)
 | 
						|
        expected_container.size = 1073741824  # 1GB in bytes
 | 
						|
        expected_container.numberOfObjects = 1500
 | 
						|
 | 
						|
        self.mock_ometa.get_by_id.return_value = expected_container
 | 
						|
 | 
						|
        # Act
 | 
						|
        result = Containers.retrieve(self.container_id)
 | 
						|
 | 
						|
        # Assert
 | 
						|
        self.assertEqual(result.size, 1073741824)
 | 
						|
        self.assertEqual(result.numberOfObjects, 1500)
 | 
						|
 | 
						|
    def test_error_handling_container_not_found(self):
 | 
						|
        """Test error handling when container not found"""
 | 
						|
        # Arrange
 | 
						|
        self.mock_ometa.get_by_id.side_effect = Exception("Container not found")
 | 
						|
 | 
						|
        # Act & Assert
 | 
						|
        with self.assertRaises(Exception) as context:
 | 
						|
            Containers.retrieve("non-existent-id")
 | 
						|
 | 
						|
        self.assertIn("Container not found", str(context.exception))
 | 
						|
 | 
						|
 | 
						|
if __name__ == "__main__":
 | 
						|
    unittest.main()
 |