| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # 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. |
| # |
| import os |
| from json.encoder import JSONEncoder |
| from typing import Any, Dict, List, Tuple |
| from unittest import TestCase |
| from unittest.mock import AsyncMock, MagicMock, call, patch |
| |
| from streampipes.client.client import StreamPipesClient, StreamPipesClientConfig |
| from streampipes.client.credential_provider import StreamPipesApiKeyCredentials |
| from streampipes.functions.broker.broker_handler import ( |
| SupportedBroker, |
| UnsupportedBrokerError, |
| ) |
| from streampipes.functions.function_handler import FunctionHandler |
| from streampipes.functions.registration import Registration |
| from streampipes.functions.streampipes_function import StreamPipesFunction |
| from streampipes.functions.utils.data_stream_generator import ( |
| RuntimeType, |
| create_data_stream, |
| ) |
| from streampipes.functions.utils.function_context import FunctionContext |
| from streampipes.model.resource.data_stream import DataStream |
| from streampipes.model.resource.function_definition import FunctionDefinition |
| |
| |
| class TestFunction(StreamPipesFunction): |
| def onServiceStarted(self, context: FunctionContext): |
| self.context = context |
| self.data: List[Dict[str, Any]] = [] |
| |
| def onEvent(self, event: Dict[str, Any], streamId: str): |
| self.data.append(event) |
| |
| def onServiceStopped(self): |
| self.stopped = True |
| |
| |
| class TestFunctionTwoStreams(StreamPipesFunction): |
| def onServiceStarted(self, context: FunctionContext): |
| self.context = context |
| self.data1: List[Dict[str, Any]] = [] |
| self.data2: List[Dict[str, Any]] = [] |
| |
| def onEvent(self, event: Dict[str, Any], streamId: str): |
| if streamId == self.requiredStreamIds()[0]: |
| self.data1.append(event) |
| else: |
| self.data2.append(event) |
| |
| def onServiceStopped(self): |
| self.stopped = True |
| |
| |
| class TestFunctionOutput(StreamPipesFunction): |
| def onServiceStarted(self, context: FunctionContext): |
| self.context = context |
| self.i = 0 |
| |
| def onEvent(self, event: Dict[str, Any], streamId: str): |
| self.add_output(self.function_definition.get_output_stream_ids()[0], event={"number": self.i}) |
| self.i += 1 |
| |
| def onServiceStopped(self): |
| self.stopped = True |
| |
| |
| class TestNatsMessage: |
| def __init__(self, data) -> None: |
| self.data = JSONEncoder().encode(data).encode() |
| |
| |
| class TestMessageIterator: |
| def __init__(self, test_data) -> None: |
| self.test_data = test_data |
| self.i = -1 |
| |
| def __aiter__(self): |
| return self |
| |
| async def __anext__(self): |
| if self.i < len(self.test_data) - 1: |
| self.i += 1 |
| return TestNatsMessage(self.test_data[self.i]) |
| else: |
| raise StopAsyncIteration |
| |
| |
| class TestKafkaMessage: |
| def __init__(self, data) -> None: |
| self.data = JSONEncoder().encode(data).encode() |
| |
| def value(self): |
| return self.data |
| |
| |
| class TestKafkaMessageContainer: |
| def __init__(self, test_data) -> None: |
| self.test_data = test_data |
| self.i = -1 |
| |
| def get_data(self, *args, **kwargs): |
| if self.i < len(self.test_data) - 1: |
| self.i += 1 |
| return TestKafkaMessage(self.test_data[self.i]) |
| else: |
| raise StopAsyncIteration |
| |
| |
| class TestFunctionHandler(TestCase): |
| def setUp(self) -> None: |
| attributes = { |
| "density": RuntimeType.FLOAT.value, |
| "temperature": RuntimeType.FLOAT.value, |
| "timestamp": RuntimeType.INTEGER.value, |
| } |
| stream_id = "urn:streampipes.apache.org:eventstream:uPDKLI" |
| |
| data_stream1 = create_data_stream("Test_NATS", attributes=attributes, stream_id=stream_id) |
| data_stream1.event_grounding.transport_protocols[0].topic_definition.actual_topic_name = "test1" |
| self.data_stream_nats: Dict[str, Any] = data_stream1.to_dict() |
| |
| data_stream2 = create_data_stream( |
| "Test_Kafka", attributes=attributes, stream_id=stream_id, broker=SupportedBroker.KAFKA |
| ) |
| data_stream2.event_grounding.transport_protocols[0].topic_definition.actual_topic_name = "test1" |
| self.data_stream_kafka: Dict[str, Any] = data_stream2.to_dict() |
| |
| data_stream1.event_grounding.transport_protocols[0].class_name = "None" |
| self.data_stream_unsupported_broker = data_stream1.to_dict() |
| |
| self.test_stream_data1 = [ |
| {"density": 10.3, "temperature": 20.5, "timestamp": 1670000001000}, |
| {"density": 13.4, "temperature": 20.4, "timestamp": 1670000002000}, |
| {"density": 12.6, "temperature": 20.7, "timestamp": 1670000003000}, |
| {"density": 9.0, "temperature": 21.0, "timestamp": 1670000004000}, |
| {"density": 10.6, "temperature": 20.4, "timestamp": 1670000005000}, |
| {"density": 11.3, "temperature": 19.5, "timestamp": 1670000006000}, |
| {"density": 15.7, "temperature": 19.3, "timestamp": 1670000007000}, |
| ] |
| self.test_stream_data2 = [ |
| {"density": 5.3, "temperature": 30.5, "timestamp": 1670000001000}, |
| {"density": 2.6, "temperature": 31.7, "timestamp": 1670000002000}, |
| {"density": 3.6, "temperature": 30.4, "timestamp": 1670000003000}, |
| {"density": 5.7, "temperature": 29.3, "timestamp": 1670000004000}, |
| {"density": 5.3, "temperature": 30.5, "timestamp": 1670000005000}, |
| {"density": 2.6, "temperature": 31.7, "timestamp": 1670000006000}, |
| {"density": 3.6, "temperature": 30.4, "timestamp": 1670000007000}, |
| ] |
| |
| @patch("streampipes.functions.broker.nats.nats_consumer.connect", autospec=True) |
| @patch("streampipes.functions.broker.NatsConsumer.get_message", autospec=True) |
| @patch("streampipes.client.client.Session", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_function_handler_nats( |
| self, server_version: MagicMock, http_session: MagicMock, get_messages: MagicMock, connection: AsyncMock |
| ): |
| http_session_mock = MagicMock() |
| http_session_mock.get.return_value.json.return_value = self.data_stream_nats |
| http_session.return_value = http_session_mock |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| |
| get_messages.return_value = TestMessageIterator(self.test_stream_data1) |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| registration = Registration() |
| test_function = TestFunction( |
| FunctionDefinition(consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"]) |
| ) |
| registration.register(test_function) |
| function_handler = FunctionHandler(registration, client) |
| function_handler.initializeFunctions() |
| |
| self.assertEqual(test_function.context.client, client) |
| self.assertDictEqual( |
| test_function.context.schema, {self.data_stream_nats["elementId"]: DataStream(**self.data_stream_nats)} |
| ) |
| self.assertListEqual(test_function.context.streams, test_function.requiredStreamIds()) |
| self.assertEqual(test_function.context.function_id, test_function.getFunctionId().id) |
| |
| self.assertListEqual(test_function.data, self.test_stream_data1) |
| self.assertTrue(test_function.stopped) |
| |
| @patch("streampipes.functions.broker.kafka.kafka_consumer.KafkaConnection", autospec=True) |
| @patch("streampipes.client.client.Session", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_function_handler_kafka(self, server_version: MagicMock, http_session: MagicMock, connection: MagicMock): |
| http_session_mock = MagicMock() |
| http_session_mock.get.return_value.json.return_value = self.data_stream_kafka |
| http_session.return_value = http_session_mock |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| |
| connection_mock = MagicMock() |
| connection_mock.poll.side_effect = TestKafkaMessageContainer(self.test_stream_data1).get_data |
| connection.return_value = connection_mock |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| registration = Registration() |
| test_function = TestFunction( |
| FunctionDefinition(consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"]) |
| ) |
| registration.register(test_function) |
| function_handler = FunctionHandler(registration, client) |
| function_handler.initializeFunctions() |
| |
| self.assertEqual(test_function.context.client, client) |
| self.assertDictEqual( |
| test_function.context.schema, {self.data_stream_kafka["elementId"]: DataStream(**self.data_stream_kafka)} |
| ) |
| self.assertListEqual(test_function.context.streams, test_function.requiredStreamIds()) |
| self.assertEqual(test_function.context.function_id, test_function.getFunctionId().id) |
| |
| self.assertListEqual(test_function.data, self.test_stream_data1) |
| self.assertTrue(test_function.stopped) |
| |
| @patch("streampipes.functions.broker.nats.nats_consumer.connect", autospec=True) |
| @patch("streampipes.client.client.Session", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_function_handler_unsupported_broker( |
| self, server_version: MagicMock, http_session: MagicMock, connection: AsyncMock |
| ): |
| http_session_mock = MagicMock() |
| http_session_mock.get.return_value.json.return_value = self.data_stream_unsupported_broker |
| http_session.return_value = http_session_mock |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| registration = Registration() |
| test_function = TestFunction( |
| FunctionDefinition(consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"]) |
| ) |
| registration.register(test_function) |
| function_handler = FunctionHandler(registration, client) |
| with self.assertRaises(UnsupportedBrokerError): |
| function_handler.initializeFunctions() |
| |
| @patch("streampipes.functions.broker.nats.nats_consumer.connect", autospec=True) |
| @patch("streampipes.functions.broker.NatsConsumer.get_message", autospec=True) |
| @patch("streampipes.endpoint.api.DataStreamEndpoint.get", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_two_streams_nats( |
| self, server_version: MagicMock, endpoint: MagicMock, nats_broker: MagicMock, *args: Tuple[AsyncMock] |
| ): |
| def get_stream(endpoint, stream_id): |
| if stream_id == "urn:streampipes.apache.org:eventstream:uPDKLI": |
| return DataStream(**self.data_stream_nats) |
| elif stream_id == "urn:streampipes.apache.org:eventstream:HHoidJ": |
| data_stream = DataStream(**self.data_stream_nats) |
| data_stream.element_id = stream_id |
| data_stream.event_grounding.transport_protocols[0].topic_definition.actual_topic_name = "test2" |
| return data_stream |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| endpoint.side_effect = get_stream |
| |
| def get_message(broker): |
| if broker.topic_name == "test1": |
| return TestMessageIterator(self.test_stream_data1) |
| elif broker.topic_name == "test2": |
| return TestMessageIterator(self.test_stream_data2) |
| |
| nats_broker.side_effect = get_message |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| registration = Registration() |
| test_function1 = TestFunction( |
| FunctionDefinition(consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"]) |
| ) |
| test_function2 = TestFunctionTwoStreams( |
| FunctionDefinition( |
| consumed_streams=[ |
| "urn:streampipes.apache.org:eventstream:uPDKLI", |
| "urn:streampipes.apache.org:eventstream:HHoidJ", |
| ] |
| ) |
| ) |
| registration.register(test_function1).register(test_function2) |
| function_handler = FunctionHandler(registration, client) |
| function_handler.initializeFunctions() |
| |
| self.assertListEqual(test_function1.data, self.test_stream_data1) |
| self.assertTrue(test_function1.stopped) |
| |
| stream_ids = test_function2.requiredStreamIds() |
| self.assertDictEqual( |
| test_function2.context.schema, |
| { |
| stream_ids[0]: get_stream(None, stream_ids[0]), |
| stream_ids[1]: get_stream(None, stream_ids[1]), |
| }, |
| ) |
| self.assertListEqual(test_function2.context.streams, test_function2.requiredStreamIds()) |
| self.assertListEqual(test_function2.data1, self.test_stream_data1) |
| self.assertListEqual(test_function2.data2, self.test_stream_data2) |
| self.assertTrue(test_function2.stopped) |
| |
| @patch("streampipes.functions.broker.nats.nats_publisher.connect", autospec=True) |
| @patch("streampipes.functions.broker.nats.nats_consumer.connect", autospec=True) |
| @patch("streampipes.functions.streampipes_function.time", autospec=True) |
| @patch("streampipes.functions.broker.NatsConsumer.get_message", autospec=True) |
| @patch("streampipes.functions.broker.NatsPublisher.publish_event", autospec=True) |
| @patch("streampipes.client.client.Session", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_function_output_stream_nats( |
| self, |
| server_version: MagicMock, |
| http_session: MagicMock, |
| pulish_event: MagicMock, |
| get_message: MagicMock, |
| time: MagicMock, |
| *args: Tuple[AsyncMock] |
| ): |
| http_session_mock = MagicMock() |
| http_session_mock.get.return_value.json.return_value = self.data_stream_nats |
| http_session.return_value = http_session_mock |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| |
| output_events = [] |
| |
| def save_event(self, event: Dict[str, Any]): |
| output_events.append(event) |
| |
| pulish_event.side_effect = save_event |
| get_message.return_value = TestMessageIterator(self.test_stream_data1) |
| time.side_effect = lambda: 0 |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| output_stream = create_data_stream("test", attributes={"number": RuntimeType.INTEGER.value}, stream_id="a1d") |
| test_function = TestFunctionOutput( |
| function_definition=FunctionDefinition( |
| consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"] |
| ).add_output_data_stream(output_stream) |
| ) |
| registration = Registration() |
| registration.register(test_function) |
| function_handler = FunctionHandler(registration, client) |
| function_handler.initializeFunctions() |
| |
| self.assertEqual(test_function.context.client, client) |
| self.assertDictEqual( |
| test_function.context.schema, {self.data_stream_nats["elementId"]: DataStream(**self.data_stream_nats)} |
| ) |
| self.assertListEqual(test_function.context.streams, test_function.requiredStreamIds()) |
| self.assertEqual(test_function.context.function_id, test_function.getFunctionId().id) |
| self.assertTrue(test_function.stopped) |
| |
| self.assertListEqual(output_events, [{"number": i, "timestamp": 0} for i in range(len(self.test_stream_data1))]) |
| |
| @patch("streampipes.functions.broker.kafka.kafka_publisher.Producer", autospec=True) |
| @patch("streampipes.functions.streampipes_function.time", autospec=True) |
| @patch("streampipes.functions.broker.kafka.kafka_consumer.KafkaConnection", autospec=True) |
| @patch("streampipes.functions.broker.KafkaPublisher.publish_event", autospec=True) |
| @patch("streampipes.client.client.Session", autospec=True) |
| @patch("streampipes.client.client.StreamPipesClient._get_server_version", autospec=True) |
| def test_function_output_stream_kafka( |
| self, |
| server_version: MagicMock, |
| http_session: MagicMock, |
| pulish_event: MagicMock, |
| connection: MagicMock, |
| time: MagicMock, |
| producer: MagicMock, |
| ): |
| os.environ["BROKER-HOST"] = "localhost" |
| os.environ["KAFKA-PORT"] = "9094" |
| http_session_mock = MagicMock() |
| http_session_mock.get.return_value.json.return_value = self.data_stream_kafka |
| http_session.return_value = http_session_mock |
| |
| server_version.return_value = {"backendVersion": "0.x.y"} |
| |
| output_events = [] |
| |
| def save_event(self, event: Dict[str, Any]): |
| output_events.append(event) |
| |
| pulish_event.side_effect = save_event |
| connection_mock = MagicMock() |
| connection_mock.poll.side_effect = TestKafkaMessageContainer(self.test_stream_data1).get_data |
| connection.return_value = connection_mock |
| time.side_effect = lambda: 0 |
| |
| client = StreamPipesClient( |
| client_config=StreamPipesClientConfig( |
| credential_provider=StreamPipesApiKeyCredentials(username="user", api_key="key"), |
| host_address="localhost", |
| ) |
| ) |
| |
| output_stream = create_data_stream( |
| "test", stream_id="a1c", attributes={"number": RuntimeType.INTEGER.value}, broker=SupportedBroker.KAFKA |
| ) |
| test_function = TestFunctionOutput( |
| function_definition=FunctionDefinition( |
| consumed_streams=["urn:streampipes.apache.org:eventstream:uPDKLI"] |
| ).add_output_data_stream(output_stream) |
| ) |
| registration = Registration() |
| registration.register(test_function) |
| function_handler = FunctionHandler(registration, client) |
| function_handler.initializeFunctions() |
| |
| producer.assert_has_calls(calls=[call({"bootstrap.servers": "localhost:9094"})]) |
| self.assertEqual(test_function.context.client, client) |
| self.assertDictEqual( |
| test_function.context.schema, {self.data_stream_kafka["elementId"]: DataStream(**self.data_stream_kafka)} |
| ) |
| self.assertListEqual(test_function.context.streams, test_function.requiredStreamIds()) |
| self.assertEqual(test_function.context.function_id, test_function.getFunctionId().id) |
| self.assertTrue(test_function.stopped) |
| |
| self.assertListEqual(output_events, [{"number": i, "timestamp": 0} for i in range(len(self.test_stream_data1))]) |
