wehub-resource-sync
bf2343b7e4
Integration Tests - MySQL + Elasticsearch / Detect Changes (push) Has been cancelled
Integration Tests - MySQL + Elasticsearch / integration-tests-mysql-elasticsearch (push) Has been cancelled
Integration Tests - PostgreSQL + Elasticsearch + Redis / Detect Changes (push) Has been cancelled
Integration Tests - PostgreSQL + Elasticsearch + Redis / integration-tests-postgres-elasticsearch-redis (push) Has been cancelled
Integration Tests - PostgreSQL + OpenSearch / Detect Changes (push) Has been cancelled
Integration Tests - PostgreSQL + OpenSearch / integration-tests-postgres-opensearch (push) Has been cancelled
Java Checkstyle / java-checkstyle (push) Has been cancelled
Maven Collate Tests / maven-collate-ci (push) Has been cancelled
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests-status (push) Has been cancelled
Publish Package to Maven Central Repository / publish-maven-packages (push) Has been cancelled
OpenMetadata Service Unit Tests / Detect Changes (push) Has been cancelled
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests (push) Has been cancelled
OpenMetadata Service Unit Tests / k8s_operator-unit-tests (push) Has been cancelled
730 lines
29 KiB
Python
730 lines
29 KiB
Python
# Copyright 2025 Collate
|
|
# Licensed under the Collate Community License, Version 1.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
Unit tests for Google Cloud Pub/Sub connector
|
|
"""
|
|
|
|
import os
|
|
import uuid
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from metadata.ingestion.source.messaging.pubsub.models import (
|
|
PubSubBigQueryConfig,
|
|
PubSubSchemaInfo,
|
|
PubSubSubscription,
|
|
PubSubTopicMetadata,
|
|
)
|
|
|
|
|
|
class TestPubSubModels:
|
|
"""Test Pub/Sub Pydantic models"""
|
|
|
|
def test_pubsub_bigquery_config_all_fields(self):
|
|
"""Test PubSubBigQueryConfig with all fields"""
|
|
config = PubSubBigQueryConfig(
|
|
table="project.dataset.table",
|
|
use_topic_schema=True,
|
|
write_metadata=True,
|
|
drop_unknown_fields=False,
|
|
)
|
|
assert config.table == "project.dataset.table"
|
|
assert config.use_topic_schema is True
|
|
assert config.write_metadata is True
|
|
assert config.drop_unknown_fields is False
|
|
|
|
def test_pubsub_bigquery_config_optional_fields(self):
|
|
"""Test PubSubBigQueryConfig with only required fields"""
|
|
config = PubSubBigQueryConfig()
|
|
assert config.table is None
|
|
assert config.use_topic_schema is None
|
|
assert config.write_metadata is None
|
|
assert config.drop_unknown_fields is None
|
|
|
|
def test_pubsub_subscription_all_fields(self):
|
|
"""Test PubSubSubscription with all fields"""
|
|
bigquery_config = PubSubBigQueryConfig(table="project.dataset.table")
|
|
subscription = PubSubSubscription(
|
|
name="test-subscription",
|
|
ack_deadline_seconds=60,
|
|
message_retention_duration=604800000.0,
|
|
dead_letter_topic="projects/test/topics/dead-letter",
|
|
push_endpoint="https://example.com/push",
|
|
filter='attributes.type = "test"',
|
|
bigquery_config=bigquery_config,
|
|
enable_exactly_once_delivery=True,
|
|
)
|
|
assert subscription.name == "test-subscription"
|
|
assert subscription.ack_deadline_seconds == 60
|
|
assert subscription.message_retention_duration == 604800000.0
|
|
assert subscription.dead_letter_topic == "projects/test/topics/dead-letter"
|
|
assert subscription.push_endpoint == "https://example.com/push"
|
|
assert subscription.filter == 'attributes.type = "test"'
|
|
assert subscription.bigquery_config.table == "project.dataset.table"
|
|
assert subscription.enable_exactly_once_delivery is True
|
|
|
|
def test_pubsub_subscription_minimal(self):
|
|
"""Test PubSubSubscription with only required fields"""
|
|
subscription = PubSubSubscription(name="minimal-subscription")
|
|
assert subscription.name == "minimal-subscription"
|
|
assert subscription.ack_deadline_seconds is None
|
|
assert subscription.bigquery_config is None
|
|
|
|
def test_pubsub_schema_info(self):
|
|
"""Test PubSubSchemaInfo model"""
|
|
schema_info = PubSubSchemaInfo(
|
|
name="test-schema",
|
|
schema_type="AVRO",
|
|
definition='{"type": "record", "name": "Test"}',
|
|
revision_id="abc123",
|
|
)
|
|
assert schema_info.name == "test-schema"
|
|
assert schema_info.schema_type == "AVRO"
|
|
assert schema_info.definition == '{"type": "record", "name": "Test"}'
|
|
assert schema_info.revision_id == "abc123"
|
|
|
|
def test_pubsub_topic_metadata_all_fields(self):
|
|
"""Test PubSubTopicMetadata with all fields"""
|
|
schema_info = PubSubSchemaInfo(name="schema", schema_type="AVRO")
|
|
subscription = PubSubSubscription(name="sub1")
|
|
|
|
metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
labels={"env": "test", "team": "data"},
|
|
message_retention_duration=604800000.0,
|
|
schema_settings=schema_info,
|
|
subscriptions=[subscription],
|
|
ordering_enabled=True,
|
|
kms_key_name="projects/test/locations/us/keyRings/ring/cryptoKeys/key",
|
|
)
|
|
assert metadata.name == "projects/test/topics/test-topic"
|
|
assert metadata.labels == {"env": "test", "team": "data"}
|
|
assert metadata.message_retention_duration == 604800000.0
|
|
assert metadata.schema_settings.name == "schema"
|
|
assert len(metadata.subscriptions) == 1
|
|
assert metadata.ordering_enabled is True
|
|
assert "keyRings" in metadata.kms_key_name
|
|
|
|
def test_pubsub_topic_metadata_minimal(self):
|
|
"""Test PubSubTopicMetadata with only required fields"""
|
|
metadata = PubSubTopicMetadata(name="projects/test/topics/minimal")
|
|
assert metadata.name == "projects/test/topics/minimal"
|
|
assert metadata.labels is None
|
|
assert metadata.subscriptions is None
|
|
assert metadata.ordering_enabled is False
|
|
|
|
|
|
class TestPubSubConnection:
|
|
"""Test Pub/Sub connection handling"""
|
|
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.set_google_credentials")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.PublisherClient")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.SubscriberClient")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.SchemaServiceClient")
|
|
def test_get_connection_with_project_id(self, mock_schema_client, mock_subscriber, mock_publisher, mock_set_creds):
|
|
"""Test get_connection with explicit project ID"""
|
|
from metadata.ingestion.source.messaging.pubsub.connection import (
|
|
PubSubClient,
|
|
get_connection,
|
|
)
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = "test-project"
|
|
mock_connection.gcpConfig = MagicMock()
|
|
mock_connection.useEmulator = False
|
|
mock_connection.hostPort = None
|
|
mock_connection.schemaRegistryEnabled = True
|
|
|
|
client = get_connection(mock_connection)
|
|
|
|
assert isinstance(client, PubSubClient)
|
|
assert client.project_id == "test-project"
|
|
mock_set_creds.assert_called_once_with(mock_connection.gcpConfig)
|
|
mock_publisher.assert_called_once()
|
|
mock_subscriber.assert_called_once()
|
|
mock_schema_client.assert_called_once()
|
|
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.set_google_credentials")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.PublisherClient")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.SubscriberClient")
|
|
def test_get_connection_without_schema_registry(self, mock_subscriber, mock_publisher, mock_set_creds):
|
|
"""Test get_connection with schema registry disabled"""
|
|
from metadata.ingestion.source.messaging.pubsub.connection import get_connection
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = "test-project"
|
|
mock_connection.gcpConfig = MagicMock()
|
|
mock_connection.useEmulator = False
|
|
mock_connection.hostPort = None
|
|
mock_connection.schemaRegistryEnabled = False
|
|
|
|
client = get_connection(mock_connection)
|
|
|
|
assert client.schema_client is None
|
|
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.set_google_credentials")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.PublisherClient")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.SubscriberClient")
|
|
def test_get_connection_with_emulator(self, mock_subscriber, mock_publisher, mock_set_creds):
|
|
"""Test get_connection with emulator enabled"""
|
|
pytest.importorskip("google.cloud.pubsub_v1", reason="google-cloud-pubsub not installed")
|
|
from metadata.ingestion.source.messaging.pubsub.connection import (
|
|
PUBSUB_EMULATOR_HOST,
|
|
get_connection,
|
|
)
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = "test-project"
|
|
mock_connection.gcpConfig = MagicMock()
|
|
mock_connection.useEmulator = True
|
|
mock_connection.hostPort = "localhost:8085"
|
|
mock_connection.schemaRegistryEnabled = False
|
|
|
|
client = get_connection(mock_connection)
|
|
|
|
assert client.project_id == "test-project"
|
|
mock_publisher.assert_called_once()
|
|
mock_subscriber.assert_called_once()
|
|
mock_set_creds.assert_not_called()
|
|
assert PUBSUB_EMULATOR_HOST not in os.environ
|
|
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.set_google_credentials")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.PublisherClient")
|
|
@patch("metadata.ingestion.source.messaging.pubsub.connection.pubsub_v1.SubscriberClient")
|
|
def test_get_connection_missing_project_id_raises(self, mock_subscriber, mock_publisher, mock_set_creds):
|
|
"""Test get_connection raises ValueError when project ID is missing"""
|
|
from metadata.ingestion.source.messaging.pubsub.connection import get_connection
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = None
|
|
mock_connection.gcpConfig = MagicMock()
|
|
mock_connection.gcpConfig.gcpConfig = None
|
|
mock_connection.useEmulator = False
|
|
mock_connection.hostPort = None
|
|
mock_connection.schemaRegistryEnabled = False
|
|
|
|
with pytest.raises(ValueError, match="Project ID is required"):
|
|
get_connection(mock_connection)
|
|
|
|
def test_get_project_id_from_connection(self):
|
|
"""Test _get_project_id extracts project ID from connection config"""
|
|
from metadata.ingestion.source.messaging.pubsub.connection import (
|
|
_get_project_id,
|
|
)
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = "explicit-project"
|
|
|
|
result = _get_project_id(mock_connection)
|
|
assert result == "explicit-project"
|
|
|
|
def test_get_project_id_from_credentials(self):
|
|
"""Test _get_project_id extracts project ID from GCP credentials"""
|
|
from metadata.generated.schema.security.credentials.gcpValues import (
|
|
GcpCredentialsValues,
|
|
SingleProjectId,
|
|
)
|
|
from metadata.ingestion.source.messaging.pubsub.connection import (
|
|
_get_project_id,
|
|
)
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = None
|
|
mock_connection.gcpConfig = MagicMock()
|
|
mock_connection.gcpConfig.gcpConfig = GcpCredentialsValues(
|
|
projectId=SingleProjectId("credentials-project"),
|
|
privateKey="fake-key",
|
|
clientEmail="test@example.iam.gserviceaccount.com",
|
|
)
|
|
|
|
result = _get_project_id(mock_connection)
|
|
assert result == "credentials-project"
|
|
|
|
def test_get_project_id_returns_none_when_missing(self):
|
|
"""Test _get_project_id returns None when project ID cannot be determined"""
|
|
from metadata.ingestion.source.messaging.pubsub.connection import (
|
|
_get_project_id,
|
|
)
|
|
|
|
mock_connection = MagicMock()
|
|
mock_connection.projectId = None
|
|
mock_connection.gcpConfig = None
|
|
|
|
result = _get_project_id(mock_connection)
|
|
assert result is None
|
|
|
|
|
|
class TestPubSubMetadataParsing:
|
|
"""Test Pub/Sub metadata parsing utilities"""
|
|
|
|
@pytest.fixture
|
|
def pubsub_source_class(self):
|
|
"""Import PubsubSource class"""
|
|
pytest.importorskip("google.cloud.pubsub_v1", reason="google-cloud-pubsub not installed")
|
|
from metadata.ingestion.source.messaging.pubsub.metadata import PubsubSource
|
|
|
|
return PubsubSource
|
|
|
|
def test_parse_retention_with_protobuf_duration(self, pubsub_source_class):
|
|
"""Test _parse_retention with protobuf Duration object"""
|
|
from google.protobuf.duration_pb2 import Duration
|
|
|
|
duration = Duration(seconds=604800, nanos=0)
|
|
result = pubsub_source_class._parse_retention(None, duration)
|
|
assert result == 604800000.0
|
|
|
|
def test_parse_retention_with_string_seconds_suffix(self, pubsub_source_class):
|
|
"""Test _parse_retention with string ending in 's'"""
|
|
result = pubsub_source_class._parse_retention(None, "604800s")
|
|
assert result == 604800000.0
|
|
|
|
def test_parse_retention_with_string_seconds_word(self, pubsub_source_class):
|
|
"""Test _parse_retention with string containing 'seconds'"""
|
|
result = pubsub_source_class._parse_retention(None, "604800 seconds")
|
|
assert result == 604800000.0
|
|
|
|
def test_parse_retention_with_none(self, pubsub_source_class):
|
|
"""Test _parse_retention with None input"""
|
|
result = pubsub_source_class._parse_retention(None, None)
|
|
assert result == 0.0
|
|
|
|
def test_parse_retention_with_invalid_string(self, pubsub_source_class):
|
|
"""Test _parse_retention with invalid string"""
|
|
result = pubsub_source_class._parse_retention(None, "invalid")
|
|
assert result == 0.0
|
|
|
|
def test_map_schema_type_avro(self, pubsub_source_class):
|
|
"""Test _map_schema_type for AVRO"""
|
|
from metadata.generated.schema.type.schema import SchemaType
|
|
|
|
result = pubsub_source_class._map_schema_type(None, "AVRO")
|
|
assert result == SchemaType.Avro
|
|
|
|
def test_map_schema_type_protobuf(self, pubsub_source_class):
|
|
"""Test _map_schema_type for PROTOCOL_BUFFER"""
|
|
from metadata.generated.schema.type.schema import SchemaType
|
|
|
|
result = pubsub_source_class._map_schema_type(None, "PROTOCOL_BUFFER")
|
|
assert result == SchemaType.Protobuf
|
|
|
|
def test_map_schema_type_unknown(self, pubsub_source_class):
|
|
"""Test _map_schema_type for unknown type"""
|
|
from metadata.generated.schema.type.schema import SchemaType
|
|
|
|
result = pubsub_source_class._map_schema_type(None, "UNKNOWN")
|
|
assert result == SchemaType.Other
|
|
|
|
|
|
class TestPubSubSourceCreation:
|
|
"""Test PubsubSource.create() method"""
|
|
|
|
def test_create_with_invalid_connection_type_raises(self):
|
|
"""Test create() raises InvalidSourceException for wrong connection type"""
|
|
from metadata.ingestion.api.steps import InvalidSourceException
|
|
from metadata.ingestion.source.messaging.pubsub.metadata import PubsubSource
|
|
|
|
config_dict = {
|
|
"type": "pubsub",
|
|
"serviceName": "test-pubsub",
|
|
"serviceConnection": {
|
|
"config": {
|
|
"type": "Kafka",
|
|
"bootstrapServers": "localhost:9092",
|
|
}
|
|
},
|
|
"sourceConfig": {
|
|
"config": {
|
|
"type": "MessagingMetadata",
|
|
}
|
|
},
|
|
}
|
|
|
|
mock_metadata = MagicMock()
|
|
|
|
with pytest.raises(InvalidSourceException):
|
|
PubsubSource.create(config_dict, mock_metadata)
|
|
|
|
|
|
class TestPubSubEdgeCases:
|
|
"""Test edge cases and error handling"""
|
|
|
|
def test_subscription_with_empty_push_config(self):
|
|
"""Test handling subscription with empty push config"""
|
|
subscription = PubSubSubscription(
|
|
name="test-sub",
|
|
push_endpoint=None,
|
|
)
|
|
assert subscription.push_endpoint is None
|
|
|
|
def test_topic_metadata_with_empty_labels(self):
|
|
"""Test handling topic with empty labels dict"""
|
|
metadata = PubSubTopicMetadata(
|
|
name="test-topic",
|
|
labels={},
|
|
)
|
|
assert metadata.labels == {}
|
|
|
|
def test_topic_metadata_with_multiple_subscriptions(self):
|
|
"""Test topic with multiple subscriptions"""
|
|
subscriptions = [PubSubSubscription(name=f"sub-{i}") for i in range(5)]
|
|
metadata = PubSubTopicMetadata(
|
|
name="test-topic",
|
|
subscriptions=subscriptions,
|
|
)
|
|
assert len(metadata.subscriptions) == 5
|
|
assert metadata.subscriptions[0].name == "sub-0"
|
|
assert metadata.subscriptions[4].name == "sub-4"
|
|
|
|
def test_bigquery_config_table_parsing(self):
|
|
"""Test BigQuery table FQN parsing"""
|
|
config = PubSubBigQueryConfig(
|
|
table="project.dataset.table",
|
|
)
|
|
parts = config.table.split(".")
|
|
assert len(parts) == 3
|
|
assert parts[0] == "project"
|
|
assert parts[1] == "dataset"
|
|
assert parts[2] == "table"
|
|
|
|
def test_schema_info_with_complex_definition(self):
|
|
"""Test schema info with complex Avro definition"""
|
|
complex_schema = """
|
|
{
|
|
"type": "record",
|
|
"name": "User",
|
|
"namespace": "com.example",
|
|
"fields": [
|
|
{"name": "id", "type": "long"},
|
|
{"name": "name", "type": "string"},
|
|
{"name": "email", "type": ["null", "string"], "default": null}
|
|
]
|
|
}
|
|
"""
|
|
schema_info = PubSubSchemaInfo(
|
|
name="user-schema",
|
|
schema_type="AVRO",
|
|
definition=complex_schema.strip(),
|
|
)
|
|
assert "record" in schema_info.definition
|
|
assert "User" in schema_info.definition
|
|
|
|
|
|
class TestPubSubTopicLineage:
|
|
"""Test Pub/Sub topic lineage functionality"""
|
|
|
|
@pytest.fixture
|
|
def mock_pubsub_source(self):
|
|
"""Create a mock PubsubSource for lineage testing"""
|
|
from metadata.ingestion.source.messaging.pubsub.metadata import PubsubSource
|
|
|
|
source = MagicMock(spec=PubsubSource)
|
|
source.metadata = MagicMock()
|
|
source.context = MagicMock()
|
|
source.context.get.return_value.messaging_service = "test-pubsub-service"
|
|
source.yield_topic_lineage = PubsubSource.yield_topic_lineage.__get__(source, PubsubSource)
|
|
return source
|
|
|
|
def test_yield_topic_lineage_no_subscriptions(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage returns nothing when there are no subscriptions"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=None,
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_empty_subscriptions(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage returns nothing when subscriptions list is empty"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_no_bigquery_config(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage skips subscriptions without BigQuery config"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
subscription = PubSubSubscription(
|
|
name="test-subscription",
|
|
bigquery_config=None,
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_bigquery_config_no_table(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage skips subscriptions with BigQuery config but no table"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
subscription = PubSubSubscription(
|
|
name="test-subscription",
|
|
bigquery_config=PubSubBigQueryConfig(table=None),
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_table_not_found(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage skips when BigQuery table is not found"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
mock_pubsub_source.metadata.get_by_name.return_value = None
|
|
|
|
subscription = PubSubSubscription(
|
|
name="test-subscription",
|
|
bigquery_config=PubSubBigQueryConfig(table="project.dataset.table"),
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_topic_not_found(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage skips when topic entity is not found"""
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
mock_table = MagicMock()
|
|
mock_table.id = "table-id-123"
|
|
|
|
mock_pubsub_source.metadata.get_by_name.return_value = None
|
|
|
|
subscription = PubSubSubscription(
|
|
name="test-subscription",
|
|
bigquery_config=PubSubBigQueryConfig(table="project.dataset.table"),
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
with (
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.search_table_from_es") as mock_search,
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.build") as mock_fqn_build,
|
|
):
|
|
mock_search.return_value = mock_table
|
|
mock_fqn_build.return_value = "test-pubsub-service.test-topic"
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
|
|
assert len(result) == 0
|
|
|
|
def test_yield_topic_lineage_success(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage successfully creates lineage"""
|
|
from metadata.generated.schema.entity.data.topic import Topic
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
table_uuid = uuid.uuid4()
|
|
topic_uuid = uuid.uuid4()
|
|
|
|
mock_table = MagicMock()
|
|
mock_table.id = table_uuid
|
|
mock_topic = MagicMock()
|
|
mock_topic.id = topic_uuid
|
|
|
|
mock_pubsub_source.metadata.get_by_name.side_effect = lambda entity, fqn: (
|
|
mock_topic if entity == Topic else None
|
|
)
|
|
|
|
subscription = PubSubSubscription(
|
|
name="bq-subscription",
|
|
bigquery_config=PubSubBigQueryConfig(table="project.dataset.events"),
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/events-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="events-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
with (
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.search_table_from_es") as mock_search,
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.build") as mock_fqn_build,
|
|
):
|
|
mock_search.return_value = mock_table
|
|
mock_fqn_build.return_value = "test-pubsub-service.events-topic"
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
|
|
assert len(result) == 1
|
|
lineage_request = result[0].right
|
|
assert lineage_request is not None
|
|
assert lineage_request.edge.fromEntity.id.root == topic_uuid
|
|
assert lineage_request.edge.fromEntity.type == "topic"
|
|
assert lineage_request.edge.toEntity.id.root == table_uuid
|
|
assert lineage_request.edge.toEntity.type == "table"
|
|
assert "bq-subscription" in lineage_request.edge.lineageDetails.description
|
|
|
|
def test_yield_topic_lineage_multiple_subscriptions(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage handles multiple BigQuery subscriptions"""
|
|
from metadata.generated.schema.entity.data.topic import Topic
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
table_uuid1 = uuid.uuid4()
|
|
table_uuid2 = uuid.uuid4()
|
|
topic_uuid = uuid.uuid4()
|
|
|
|
mock_table1 = MagicMock()
|
|
mock_table1.id = table_uuid1
|
|
mock_table2 = MagicMock()
|
|
mock_table2.id = table_uuid2
|
|
mock_topic = MagicMock()
|
|
mock_topic.id = topic_uuid
|
|
|
|
mock_pubsub_source.metadata.get_by_name.side_effect = lambda entity, fqn: (
|
|
mock_topic if entity == Topic else None
|
|
)
|
|
|
|
search_call_count = {"count": 0}
|
|
|
|
def search_side_effect(**kwargs):
|
|
search_call_count["count"] += 1
|
|
return mock_table1 if search_call_count["count"] == 1 else mock_table2
|
|
|
|
subscriptions = [
|
|
PubSubSubscription(
|
|
name="bq-sub-1",
|
|
bigquery_config=PubSubBigQueryConfig(table="project.dataset.table1"),
|
|
),
|
|
PubSubSubscription(
|
|
name="push-sub",
|
|
push_endpoint="https://example.com/push",
|
|
),
|
|
PubSubSubscription(
|
|
name="bq-sub-2",
|
|
bigquery_config=PubSubBigQueryConfig(table="project.dataset.table2"),
|
|
),
|
|
]
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/multi-topic",
|
|
subscriptions=subscriptions,
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="multi-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
with (
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.search_table_from_es") as mock_search,
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.build") as mock_fqn_build,
|
|
):
|
|
mock_search.side_effect = search_side_effect
|
|
mock_fqn_build.return_value = "test-pubsub-service.multi-topic"
|
|
result = list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
|
|
assert len(result) == 2
|
|
assert result[0].right.edge.toEntity.id.root == table_uuid1
|
|
assert result[1].right.edge.toEntity.id.root == table_uuid2
|
|
|
|
def test_yield_topic_lineage_table_fqn_parsing(self, mock_pubsub_source):
|
|
"""Test yield_topic_lineage correctly parses BigQuery table FQN"""
|
|
from metadata.generated.schema.entity.data.topic import Topic
|
|
from metadata.ingestion.source.messaging.messaging_service import (
|
|
BrokerTopicDetails,
|
|
)
|
|
|
|
mock_table = MagicMock()
|
|
mock_table.id = "table-id"
|
|
mock_topic = MagicMock()
|
|
mock_topic.id = "topic-id"
|
|
|
|
mock_pubsub_source.metadata.get_by_name.side_effect = lambda entity, fqn: (
|
|
mock_topic if entity == Topic else None
|
|
)
|
|
|
|
subscription = PubSubSubscription(
|
|
name="bq-sub",
|
|
bigquery_config=PubSubBigQueryConfig(table="my-project.my_dataset.my_table"),
|
|
)
|
|
topic_metadata = PubSubTopicMetadata(
|
|
name="projects/test/topics/test-topic",
|
|
subscriptions=[subscription],
|
|
)
|
|
topic_details = BrokerTopicDetails(
|
|
topic_name="test-topic",
|
|
topic_metadata=topic_metadata,
|
|
)
|
|
|
|
with (
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.search_table_from_es") as mock_search,
|
|
patch("metadata.ingestion.source.messaging.pubsub.metadata.fqn.build") as mock_fqn_build,
|
|
):
|
|
mock_search.return_value = mock_table
|
|
mock_fqn_build.return_value = "test-pubsub-service.test-topic"
|
|
list(mock_pubsub_source.yield_topic_lineage(topic_details))
|
|
|
|
mock_search.assert_called_once_with(
|
|
metadata=mock_pubsub_source.metadata,
|
|
database_name="my-project",
|
|
schema_name="my_dataset",
|
|
service_name=None,
|
|
table_name="my_table",
|
|
)
|