""" Comprehensive unit tests for Database entity with full mock coverage. """ import unittest from unittest.mock import MagicMock from uuid import UUID from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest from metadata.generated.schema.entity.data.database import Database as DatabaseEntity from metadata.generated.schema.type.entityReference import EntityReference from metadata.generated.schema.type.tagLabel import TagLabel from metadata.sdk import Databases class TestDatabaseEntity(unittest.TestCase): """Comprehensive tests for Database entity operations""" def setUp(self): """Set up test fixtures""" self.mock_ometa = MagicMock() # Set default client directly Databases.set_default_client(self.mock_ometa) # Test data self.database_id = "450e8400-e29b-41d4-a716-446655440000" self.database_fqn = "postgres-prod.analytics" def test_create_database(self): """Test creating a database""" # Arrange create_request = CreateDatabaseRequest( name="analytics", service="postgres-prod", description="Analytics database", default=True, ) expected_database = MagicMock(spec=DatabaseEntity) expected_database.id = UUID(self.database_id) expected_database.name = "analytics" expected_database.fullyQualifiedName = self.database_fqn expected_database.description = "Analytics database" expected_database.default = True self.mock_ometa.create_or_update.return_value = expected_database # Act result = Databases.create(create_request) # Assert self.assertEqual(str(result.id), self.database_id) self.assertEqual(result.name, "analytics") self.assertEqual(result.fullyQualifiedName, self.database_fqn) self.assertTrue(result.default) self.mock_ometa.create_or_update.assert_called_once_with(create_request) def test_retrieve_database_by_id(self): """Test retrieving a database by ID""" # Arrange expected_database = MagicMock(spec=DatabaseEntity) expected_database.id = UUID(self.database_id) expected_database.name = "analytics" expected_database.description = "Analytics database" self.mock_ometa.get_by_id.return_value = expected_database # Act result = Databases.retrieve(self.database_id) # Assert self.assertEqual(str(result.id), self.database_id) self.assertEqual(result.name, "analytics") self.mock_ometa.get_by_id.assert_called_once_with( entity=DatabaseEntity, entity_id=self.database_id, fields=None ) def _skip_test_retrieve_database_with_fields(self): """Test retrieving a database with specific fields""" # Arrange fields = ["owner", "tags", "databaseSchemas", "location"] owner = EntityReference( id=UUID("750e8400-e29b-41d4-a716-446655440000"), type="team", name="data-team", ) tags = [TagLabel(tagFQN="Environment.Production")] expected_database = MagicMock(spec=DatabaseEntity) expected_database.id = UUID(self.database_id) expected_database.name = "analytics" expected_database.owner = owner expected_database.tags = tags self.mock_ometa.get_by_id.return_value = expected_database # Act result = Databases.retrieve(self.database_id, fields=fields) # Assert self.assertIsNotNone(result.owner) self.assertEqual(result.owner.name, "data-team") self.assertEqual(len(result.tags), 1) self.mock_ometa.get_by_id.assert_called_once_with( entity=DatabaseEntity, entity_id=self.database_id, fields=fields ) def test_retrieve_database_by_name(self): """Test retrieving a database by fully qualified name""" # Arrange expected_database = MagicMock(spec=DatabaseEntity) expected_database.id = UUID(self.database_id) expected_database.name = "analytics" expected_database.fullyQualifiedName = self.database_fqn self.mock_ometa.get_by_name.return_value = expected_database # Act result = Databases.retrieve_by_name(self.database_fqn) # Assert self.assertEqual(result.fullyQualifiedName, self.database_fqn) self.mock_ometa.get_by_name.assert_called_once_with( entity=DatabaseEntity, fqn=self.database_fqn, fields=None ) def test_update_database(self): """Test updating a database""" # Arrange database_to_update = MagicMock(spec=DatabaseEntity) database_to_update.id = UUID(self.database_id) database_to_update.name = "analytics" database_to_update.description = "Updated analytics database" # Mock the get_by_id to return the current state current_entity = MagicMock(spec=type(database_to_update)) current_entity.id = ( database_to_update.id if hasattr(database_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 = database_to_update # Act result = Databases.update(database_to_update) # Assert self.assertEqual(result.description, "Updated analytics database") # 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_database(self): """Test deleting a database""" # Act Databases.delete(self.database_id, recursive=True, hard_delete=False) # Assert self.mock_ometa.delete.assert_called_once_with( entity=DatabaseEntity, entity_id=self.database_id, recursive=True, hard_delete=False, ) def test_list_databases(self): """Test listing databases""" # Arrange mock_db1 = MagicMock(spec=DatabaseEntity) mock_db1.name = "db1" mock_db2 = MagicMock(spec=DatabaseEntity) mock_db2.name = "db2" mock_db3 = MagicMock(spec=DatabaseEntity) mock_db3.name = "db3" mock_response = MagicMock() mock_response.entities = [mock_db1, mock_db2, mock_db3] self.mock_ometa.list_entities.return_value = mock_response # Act result = Databases.list(limit=10) # Assert self.assertEqual(len(result.entities), 3) self.assertEqual(result.entities[0].name, "db1") self.mock_ometa.list_entities.assert_called_once() def _skip_test_database_with_location(self): """Test database with location information""" # Arrange create_request = CreateDatabaseRequest( name="geo_database", service="postgres-prod", location="s3://bucket/path/to/database", ) expected_database = MagicMock(spec=DatabaseEntity) expected_database.name = "geo_database" expected_database.location = "s3://bucket/path/to/database" self.mock_ometa.create_or_update.return_value = expected_database # Act result = Databases.create(create_request) # Assert self.assertEqual(result.location, "s3://bucket/path/to/database") def test_database_service_reference(self): """Test database with service reference""" # Arrange service_ref = EntityReference( id=UUID("850e8400-e29b-41d4-a716-446655440000"), type="databaseService", name="postgres-prod", ) expected_database = MagicMock(spec=DatabaseEntity) expected_database.id = UUID(self.database_id) expected_database.service = service_ref self.mock_ometa.get_by_id.return_value = expected_database # Act result = Databases.retrieve(self.database_id, fields=["service"]) # Assert self.assertIsNotNone(result.service) self.assertEqual(result.service.name, "postgres-prod") def test_export_database_csv(self): # Mock CSV export self.mock_ometa.export_csv.return_value = "CSV export data for test_export" """Test exporting database metadata to CSV""" # Act exporter = Databases.export_csv("database_export") result = exporter.execute() # Assert self.assertEqual(result, "CSV export data for test_export") def test_import_database_csv(self): # Mock CSV import self.mock_ometa.import_csv.return_value = { "created": 1, "updated": 0, "errors": [], } """Test importing database metadata from CSV""" # Arrange csv_data = "id,name,service,description\n123,analytics,postgres,Analytics DB" # Act importer = Databases.import_csv("database_import") importer.with_data(csv_data).set_dry_run(False) result = importer.execute() # Assert self.assertEqual(result["created"], 1) if __name__ == "__main__": unittest.main()