OpenMetadata/ingestion/tests/unit/sdk/test_query_entity.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

130 lines
4.5 KiB
Python

"""
Comprehensive unit tests for Query entity.
"""
import unittest
from unittest.mock import MagicMock
from uuid import UUID
from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
from metadata.generated.schema.entity.data.query import Query as QueryEntity
from metadata.sdk import Queries
class TestQueryEntity(unittest.TestCase):
"""Comprehensive tests for Query entity operations"""
def setUp(self):
"""Set up test fixtures"""
self.mock_ometa = MagicMock()
Queries.set_default_client(self.mock_ometa)
self.entity_id = "550e8400-e29b-41d4-a716-446655440000"
self.entity_fqn = "service.query.test_query"
def test_create_query(self):
"""Test creating a query"""
create_request = MagicMock(spec=CreateQueryRequest)
create_request.name = "test_query"
create_request.displayName = "Test Query"
create_request.description = "Test query for unit tests"
expected_entity = MagicMock(spec=QueryEntity)
expected_entity.id = UUID(self.entity_id)
expected_entity.name = "test_query"
self.mock_ometa.create_or_update.return_value = expected_entity
result = Queries.create(create_request)
self.assertEqual(str(result.id), self.entity_id)
self.assertEqual(result.name, "test_query")
self.mock_ometa.create_or_update.assert_called_once_with(create_request)
def test_retrieve_query_by_id(self):
"""Test retrieving a query by ID"""
expected_entity = MagicMock(spec=QueryEntity)
expected_entity.id = UUID(self.entity_id)
expected_entity.name = "test_query"
self.mock_ometa.get_by_id.return_value = expected_entity
result = Queries.retrieve(self.entity_id)
self.assertEqual(str(result.id), self.entity_id)
self.mock_ometa.get_by_id.assert_called_once_with(
entity=QueryEntity, entity_id=self.entity_id, fields=None
)
def test_retrieve_query_by_name(self):
"""Test retrieving a query by name"""
expected_entity = MagicMock(spec=QueryEntity)
expected_entity.fullyQualifiedName = self.entity_fqn
self.mock_ometa.get_by_name.return_value = expected_entity
result = Queries.retrieve_by_name(self.entity_fqn)
self.assertEqual(result.fullyQualifiedName, self.entity_fqn)
self.mock_ometa.get_by_name.assert_called_once_with(
entity=QueryEntity, fqn=self.entity_fqn, fields=None
)
def test_update_query(self):
"""Test updating a query"""
entity_to_update = MagicMock(spec=QueryEntity)
entity_to_update.id = UUID(self.entity_id)
entity_to_update.description = "Updated description"
# Mock the get_by_id to return the current state
current_entity = MagicMock(spec=type(entity_to_update))
current_entity.id = (
entity_to_update.id
if hasattr(entity_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 = entity_to_update
result = Queries.update(entity_to_update)
self.assertEqual(result.description, "Updated description")
# 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_query(self):
"""Test deleting a query"""
Queries.delete(self.entity_id, recursive=True, hard_delete=False)
self.mock_ometa.delete.assert_called_once_with(
entity=QueryEntity,
entity_id=self.entity_id,
recursive=True,
hard_delete=False,
)
def test_list_queries(self):
"""Test listing querys"""
mock_entity1 = MagicMock(spec=QueryEntity)
mock_entity1.name = "entity1"
mock_entity2 = MagicMock(spec=QueryEntity)
mock_entity2.name = "entity2"
mock_response = MagicMock()
mock_response.entities = [mock_entity1, mock_entity2]
self.mock_ometa.list_entities.return_value = mock_response
result = Queries.list(limit=10)
self.assertEqual(len(result.entities), 2)
self.assertEqual(result.entities[0].name, "entity1")
self.mock_ometa.list_entities.assert_called_once()
if __name__ == "__main__":
unittest.main()