This is an automated email from the ASF dual-hosted git repository.
zehnder pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/streampipes.git
The following commit(s) were added to refs/heads/dev by this push:
new d33800790 Python Functions: Distinction between consumer and publisher
(#1523)
d33800790 is described below
commit d3380079092b64d2b5140a336b900c411e032d92
Author: Sven Oehler <[email protected]>
AuthorDate: Thu May 4 21:54:08 2023 +0200
Python Functions: Distinction between consumer and publisher (#1523)
* Add consumer and publisher
* Configured Kafka port
* Enable output streams for Kafka
* Fix tests
* Improve broker tests
* Further improve broker tests
---
.../docs/getting-started/first-steps.md | 12 +-
...ive-data-from-the-streampipes-data-stream.ipynb | 21 +-
...ine-learning-on-a-streampipes-data-stream.ipynb | 18 +-
.../streampipes/function_zoo/river_function.py | 8 +-
.../streampipes/functions/broker/__init__.py | 19 +-
.../streampipes/functions/broker/broker.py | 48 +---
.../streampipes/functions/broker/broker_handler.py | 51 +++-
.../streampipes/functions/broker/consumer.py | 65 +++++
.../functions/broker/{ => kafka}/__init__.py | 13 -
.../{kafka_broker.py => kafka/kafka_consumer.py} | 38 +--
.../broker/{ => kafka}/kafka_message_fetcher.py | 4 +-
.../{nats_broker.py => kafka/kafka_publisher.py} | 45 +--
.../functions/broker/{ => nats}/__init__.py | 13 -
.../{nats_broker.py => nats/nats_consumer.py} | 31 +--
.../{nats_broker.py => nats/nats_publisher.py} | 34 +--
.../functions/broker/output_collector.py | 14 +-
.../functions/broker/{__init__.py => publisher.py} | 35 ++-
.../streampipes/functions/function_handler.py | 5 +-
.../functions/utils/async_iter_handler.py | 2 +-
.../functions/utils/data_stream_context.py | 8 +-
.../functions/utils/data_stream_generator.py | 34 ++-
.../tests/client/test_endpoint.py | 6 +-
.../tests/functions/test_function_handler.py | 306 ++++++++++++---------
.../tests/functions/test_river_function.py | 24 +-
24 files changed, 473 insertions(+), 381 deletions(-)
diff --git a/streampipes-client-python/docs/getting-started/first-steps.md
b/streampipes-client-python/docs/getting-started/first-steps.md
index 8c75d7636..a8f0288b4 100644
--- a/streampipes-client-python/docs/getting-started/first-steps.md
+++ b/streampipes-client-python/docs/getting-started/first-steps.md
@@ -67,7 +67,17 @@ Once all services are started, you can access StreamPipes
via `http://localhost`
#### Setup StreamPipes with Kafka as message broker
Alternatively, you can use `docker-compose.yml` to start StreamPipes with
Kafka as messaging layer.
-Therefore, you onyl need to execute the following command:
+When running locally we have to modify `services.kafka.environment` and add
the ports to `services.kafka.ports`:
+```yaml
+environment:
+ KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,OUTSIDE:PLAINTEXT
+ KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://:9092,OUTSIDE://localhost:9094
+ KAFKA_LISTENERS: PLAINTEXT://:9092,OUTSIDE://:9094
+...
+ports:
+ - 9094:9094
+```
+Then, you need to execute the following command:
```bash
docker-compose -f docker-compose.yml up -d
```
diff --git
a/streampipes-client-python/docs/tutorials/3-getting-live-data-from-the-streampipes-data-stream.ipynb
b/streampipes-client-python/docs/tutorials/3-getting-live-data-from-the-streampipes-data-stream.ipynb
index 0a2187486..5d2dcef89 100644
---
a/streampipes-client-python/docs/tutorials/3-getting-live-data-from-the-streampipes-data-stream.ipynb
+++
b/streampipes-client-python/docs/tutorials/3-getting-live-data-from-the-streampipes-data-stream.ipynb
@@ -34,14 +34,14 @@
{
"cell_type": "code",
"execution_count": null,
+ "metadata": {
+ "collapsed": false
+ },
"outputs": [],
"source": [
"# You can install all required libraries for this tutorial with the
following command\n",
"%pip install matplotlib ipython streampipes"
- ],
- "metadata": {
- "collapsed": false
- }
+ ]
},
{
"cell_type": "code",
@@ -53,7 +53,10 @@
"\n",
"os.environ[\"USER\"] = \"[email protected]\"\n",
"os.environ[\"API-KEY\"] = \"XXX\"\n",
- "os.environ[\"BROKER-HOST\"] = \"localhost\" # Use this if you work
locally"
+ "\n",
+ "# Use this if you work locally:\n",
+ "os.environ[\"BROKER-HOST\"] = \"localhost\" \n",
+ "os.environ[\"KAFKA-PORT\"] = \"9094\" # When using Kafka as message
broker"
]
},
{
@@ -591,12 +594,12 @@
},
{
"cell_type": "markdown",
- "source": [
- "Want to see more exciting use cases you can achieve with StreamPipes
functions in Python? Then don't hesitate and jump to our [next
tutorial](../4-using-online-machine-learning-on-a-streampipes-data-stream) on
applying online machine learning algorithms to StreamPipes data streams with
[River](https://riverml.xyz)."
- ],
"metadata": {
"collapsed": false
- }
+ },
+ "source": [
+ "Want to see more exciting use cases you can achieve with StreamPipes
functions in Python? Then don't hesitate and jump to our [next
tutorial](../4-using-online-machine-learning-on-a-streampipes-data-stream) on
applying online machine learning algorithms to StreamPipes data streams with
[River](https://riverml.xyz)."
+ ]
},
{
"attachments": {},
diff --git
a/streampipes-client-python/docs/tutorials/4-using-online-machine-learning-on-a-streampipes-data-stream.ipynb
b/streampipes-client-python/docs/tutorials/4-using-online-machine-learning-on-a-streampipes-data-stream.ipynb
index 8e6673ab4..780464d3d 100644
---
a/streampipes-client-python/docs/tutorials/4-using-online-machine-learning-on-a-streampipes-data-stream.ipynb
+++
b/streampipes-client-python/docs/tutorials/4-using-online-machine-learning-on-a-streampipes-data-stream.ipynb
@@ -23,14 +23,14 @@
{
"cell_type": "code",
"execution_count": null,
+ "metadata": {
+ "collapsed": false
+ },
"outputs": [],
"source": [
"# you can install all required dependecies for this tutorial by executing
the following command\n",
"%pip install river streampipes"
- ],
- "metadata": {
- "collapsed": false
- }
+ ]
},
{
"cell_type": "code",
@@ -42,7 +42,10 @@
"\n",
"os.environ[\"USER\"] = \"[email protected]\"\n",
"os.environ[\"API-KEY\"] = \"XXX\"\n",
- "os.environ[\"BROKER-HOST\"] = \"localhost\""
+ "\n",
+ "# Use this if you work locally:\n",
+ "os.environ[\"BROKER-HOST\"] = \"localhost\" \n",
+ "os.environ[\"KAFKA-PORT\"] = \"9094\" # When using Kafka as message
broker"
]
},
{
@@ -324,7 +327,8 @@
"metadata": {},
"outputs": [],
"source": [
- "from river import cluster, compose, preprocessing, tree\n",
+ "import pickle\n",
+ "from river import compose, tree\n",
"from streampipes.function_zoo.river_function import OnlineML\n",
"from streampipes.functions.utils.data_stream_generator import
RuntimeType\n",
"\n",
@@ -393,7 +397,7 @@
"outputs": [],
"source": [
"import pickle\n",
- "from river import cluster, compose, preprocessing, tree\n",
+ "from river import compose, tree\n",
"from streampipes.function_zoo.river_function import OnlineML\n",
"from streampipes.functions.utils.data_stream_generator import
RuntimeType\n",
"\n",
diff --git
a/streampipes-client-python/streampipes/function_zoo/river_function.py
b/streampipes-client-python/streampipes/function_zoo/river_function.py
index b43b40d0c..b42a7dea9 100644
--- a/streampipes-client-python/streampipes/function_zoo/river_function.py
+++ b/streampipes-client-python/streampipes/function_zoo/river_function.py
@@ -17,6 +17,7 @@
from typing import Any, Callable, Dict, List, Optional
from streampipes.client.client import StreamPipesClient
+from streampipes.functions.broker.broker_handler import get_broker_description
from streampipes.functions.function_handler import FunctionHandler
from streampipes.functions.registration import Registration
from streampipes.functions.streampipes_function import StreamPipesFunction
@@ -187,7 +188,12 @@ class OnlineML:
attributes["truth"] = prediction_type
if target_label is None:
raise ValueError("You must define a target attribute for a
supervised model.")
- output_stream = create_data_stream("prediction", attributes)
+
+ output_stream = create_data_stream(
+ name="prediction",
+ attributes=attributes,
+
broker=get_broker_description(client.dataStreamApi.get(stream_ids[0])), #
type: ignore
+ )
function_definition =
FunctionDefinition().add_output_data_stream(output_stream)
self.sp_function = RiverFunction(
function_definition, stream_ids, model, supervised, target_label,
on_start, on_event, on_stop
diff --git a/streampipes-client-python/streampipes/functions/broker/__init__.py
b/streampipes-client-python/streampipes/functions/broker/__init__.py
index 32d35fba3..687b55b36 100644
--- a/streampipes-client-python/streampipes/functions/broker/__init__.py
+++ b/streampipes-client-python/streampipes/functions/broker/__init__.py
@@ -15,15 +15,26 @@
# limitations under the License.
#
from .broker import Broker
-from .kafka_broker import KafkaBroker
-from .nats_broker import NatsBroker
+from .consumer import Consumer
+from .publisher import Publisher
+
+# isort: split
+
+from .kafka.kafka_consumer import KafkaConsumer
+from .kafka.kafka_publisher import KafkaPublisher
+from .nats.nats_consumer import NatsConsumer
+from .nats.nats_publisher import NatsPublisher
from .broker_handler import SupportedBroker, get_broker # isort: skip
__all__ = [
"Broker",
- "KafkaBroker",
- "NatsBroker",
+ "Consumer",
+ "Publisher",
"SupportedBroker",
"get_broker",
+ "KafkaConsumer",
+ "KafkaPublisher",
+ "NatsConsumer",
+ "NatsPublisher",
]
diff --git a/streampipes-client-python/streampipes/functions/broker/broker.py
b/streampipes-client-python/streampipes/functions/broker/broker.py
index 7697bda11..62f645356 100644
--- a/streampipes-client-python/streampipes/functions/broker/broker.py
+++ b/streampipes-client-python/streampipes/functions/broker/broker.py
@@ -16,15 +16,14 @@
#
import os
from abc import ABC, abstractmethod
-from typing import Any, AsyncIterator, Dict
from streampipes.model.resource.data_stream import DataStream
class Broker(ABC):
- """Abstract implementation of a broker.
+ """Abstract implementation of a broker for consumer and publisher.
- A broker allows both to subscribe to a data stream and to publish events
to a data stream.
+ It contains the basic logic to connect to a data stream.
"""
async def connect(self, data_stream: DataStream) -> None:
@@ -43,12 +42,15 @@ class Broker(ABC):
transport_protocol = data_stream.event_grounding.transport_protocols[0]
self.topic_name = transport_protocol.topic_definition.actual_topic_name
hostname = transport_protocol.broker_hostname
+ port = transport_protocol.port
if "BROKER-HOST" in os.environ.keys():
hostname = os.environ["BROKER-HOST"]
- await self._makeConnection(hostname, transport_protocol.port)
+ if "Kafka" in transport_protocol.class_name and "KAFKA-PORT" in
os.environ.keys():
+ port = int(os.environ["KAFKA-PORT"])
+ await self._make_connection(hostname, port)
@abstractmethod
- async def _makeConnection(self, hostname: str, port: int) -> None:
+ async def _make_connection(self, hostname: str, port: int) -> None:
"""Helper function to connect to a server.
Parameters
@@ -64,31 +66,6 @@ class Broker(ABC):
"""
raise NotImplementedError # pragma: no cover
- @abstractmethod
- async def createSubscription(self) -> None:
- """Creates a subscription to a data stream.
-
- Returns
- -------
- None
- """
- raise NotImplementedError # pragma: no cover
-
- @abstractmethod
- async def publish_event(self, event: Dict[str, Any]) -> None:
- """Publish an event to a connected data stream.
-
- Parameters
- ----------
- event: Dict[str, Any]
- The event to be published.
-
- Returns
- -------
- None
- """
- raise NotImplementedError # pragma: no cover
-
@abstractmethod
async def disconnect(self) -> None:
"""Closes the connection to the server.
@@ -98,14 +75,3 @@ class Broker(ABC):
None
"""
raise NotImplementedError # pragma: no cover
-
- @abstractmethod
- def get_message(self) -> AsyncIterator:
- """Get the published messages of the subscription.
-
- Returns
- -------
- iterator: AsyncIterator
- An async iterator for the messages.
- """
- raise NotImplementedError # pragma: no cover
diff --git
a/streampipes-client-python/streampipes/functions/broker/broker_handler.py
b/streampipes-client-python/streampipes/functions/broker/broker_handler.py
index 6bf070486..ae354fe7e 100644
--- a/streampipes-client-python/streampipes/functions/broker/broker_handler.py
+++ b/streampipes-client-python/streampipes/functions/broker/broker_handler.py
@@ -16,7 +16,13 @@
#
from enum import Enum
-from streampipes.functions.broker import Broker, KafkaBroker, NatsBroker
+from streampipes.functions.broker import (
+ Broker,
+ KafkaConsumer,
+ KafkaPublisher,
+ NatsConsumer,
+ NatsPublisher,
+)
from streampipes.model.resource.data_stream import DataStream
@@ -31,11 +37,13 @@ class SupportedBroker(Enum):
class UnsupportedBrokerError(Exception):
"""Exception if a broker isn't implemented yet."""
- def __init__(self, message):
- super().__init__(message)
+ def __init__(self, broker_name: str):
+ super().__init__(f'The python client doesn\'t support the broker:
"{broker_name}" yet')
-def get_broker(data_stream: DataStream) -> Broker: # TODO implementation for
more transport_protocols
+def get_broker(
+ data_stream: DataStream, is_publisher: bool = False
+) -> Broker: # TODO implementation for more transport_protocols
"""Derive the broker for the given data stream.
Parameters
@@ -55,8 +63,37 @@ def get_broker(data_stream: DataStream) -> Broker: # TODO
implementation for mo
"""
broker_name = data_stream.event_grounding.transport_protocols[0].class_name
if SupportedBroker.NATS.value in broker_name:
- return NatsBroker()
+ if is_publisher:
+ return NatsPublisher()
+ return NatsConsumer()
elif SupportedBroker.KAFKA.value in broker_name:
- return KafkaBroker()
+ if is_publisher:
+ return KafkaPublisher()
+ return KafkaConsumer()
else:
- raise UnsupportedBrokerError(f'The python client doesn\'t include the
broker: "{broker_name}" yet')
+ raise UnsupportedBrokerError(broker_name)
+
+
+def get_broker_description(data_stream: DataStream) -> SupportedBroker:
+ """Derive the decription of the broker for the given data stream.
+
+ Parameters
+ ----------
+ data_stream: DataStream
+ Data stream instance from which the broker is inferred
+
+ Returns
+ -------
+ broker: SupportedBroker
+ The corresponding broker description derived from data stream.
+
+ Raises
+ ------
+ UnsupportedBrokerError
+ Is raised when the given data stream belongs to a broker that is
currently not supported by StreamPipes Python.
+ """
+ broker_name = data_stream.event_grounding.transport_protocols[0].class_name
+ for b in SupportedBroker:
+ if b.value in broker_name:
+ return b
+ raise UnsupportedBrokerError(broker_name)
diff --git a/streampipes-client-python/streampipes/functions/broker/consumer.py
b/streampipes-client-python/streampipes/functions/broker/consumer.py
new file mode 100644
index 000000000..03e988504
--- /dev/null
+++ b/streampipes-client-python/streampipes/functions/broker/consumer.py
@@ -0,0 +1,65 @@
+#
+# 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.
+#
+from abc import abstractmethod
+from typing import AsyncIterator
+
+from streampipes.functions.broker import Broker
+from streampipes.model.resource.data_stream import DataStream
+
+
+class Consumer(Broker):
+ """Abstract implementation a consumer for a broker.
+
+ A consumer allows to subscribe to a data stream.
+ """
+
+ async def connect(self, data_stream: DataStream) -> None:
+ """Connects to the broker running in StreamPipes and creates a
subscription.
+
+ Parameters
+ ----------
+ data_stream: DataStream
+ Contains the meta information (resources) for a data stream.
+
+ Returns
+ -------
+ None
+ """
+ await super().connect(data_stream)
+ await self._create_subscription()
+
+ @abstractmethod
+ async def _create_subscription(self) -> None:
+ """Creates a subscription to a data stream.
+
+ Returns
+ -------
+ None
+
+ """
+ raise NotImplementedError # pragma: no cover
+
+ @abstractmethod
+ def get_message(self) -> AsyncIterator:
+ """Get the published messages of the subscription.
+
+ Returns
+ -------
+ iterator: AsyncIterator
+ An async iterator for the messages.
+ """
+ raise NotImplementedError # pragma: no cover
diff --git a/streampipes-client-python/streampipes/functions/broker/__init__.py
b/streampipes-client-python/streampipes/functions/broker/kafka/__init__.py
similarity index 73%
copy from streampipes-client-python/streampipes/functions/broker/__init__.py
copy to streampipes-client-python/streampipes/functions/broker/kafka/__init__.py
index 32d35fba3..cce3acad3 100644
--- a/streampipes-client-python/streampipes/functions/broker/__init__.py
+++ b/streampipes-client-python/streampipes/functions/broker/kafka/__init__.py
@@ -14,16 +14,3 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from .broker import Broker
-from .kafka_broker import KafkaBroker
-from .nats_broker import NatsBroker
-
-from .broker_handler import SupportedBroker, get_broker # isort: skip
-
-__all__ = [
- "Broker",
- "KafkaBroker",
- "NatsBroker",
- "SupportedBroker",
- "get_broker",
-]
diff --git
a/streampipes-client-python/streampipes/functions/broker/kafka_broker.py
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_consumer.py
similarity index 69%
rename from
streampipes-client-python/streampipes/functions/broker/kafka_broker.py
rename to
streampipes-client-python/streampipes/functions/broker/kafka/kafka_consumer.py
index 9bc1337a2..bcf5215c0 100644
--- a/streampipes-client-python/streampipes/functions/broker/kafka_broker.py
+++
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_consumer.py
@@ -16,20 +16,20 @@
#
import logging
-from typing import Any, AsyncIterator, Dict
+from typing import AsyncIterator
-from confluent_kafka import Consumer # type: ignore
-from streampipes.functions.broker.broker import Broker
-from streampipes.functions.broker.kafka_message_fetcher import
KafkaMessageFetcher
+from confluent_kafka import Consumer as KafkaConnection # type: ignore
+from streampipes.functions.broker import Consumer
+from streampipes.functions.broker.kafka.kafka_message_fetcher import
KafkaMessageFetcher
from streampipes.model.common import random_letters
logger = logging.getLogger(__name__)
-class KafkaBroker(Broker):
- """Implementation of the broker for Kafka"""
+class KafkaConsumer(Consumer):
+ """Implementation of a consumer for Kafka"""
- async def _makeConnection(self, hostname: str, port: int) -> None:
+ async def _make_connection(self, hostname: str, port: int) -> None:
"""Helper function to connect to a server.
Parameters
@@ -45,11 +45,12 @@ class KafkaBroker(Broker):
-------
None
"""
- self.kafka_consumer = Consumer(
+ self.kafka_consumer = KafkaConnection(
{"bootstrap.servers": f"{hostname}:{port}", "group.id":
random_letters(6), "auto.offset.reset": "latest"}
)
+ logger.info(f"Connecting to Kafka at {hostname}:{port}")
- async def createSubscription(self) -> None:
+ async def _create_subscription(self) -> None:
"""Creates a subscription to a data stream.
Returns
@@ -57,23 +58,7 @@ class KafkaBroker(Broker):
None
"""
self.kafka_consumer.subscribe([self.topic_name])
-
- logger.info(f"Subscribed to stream: {self.stream_id}")
-
- async def publish_event(self, event: Dict[str, Any]):
- """Publish an event to a connected data stream.
-
- Parameters
- ----------
- event: Dict[str, Any]
- The event to be published.
-
- Returns
- -------
- None
- """
-
- # await self.publish(subject=self.topic_name,
payload=json.dumps(event).encode("utf-8"))
+ logger.info(f"Subscribing to stream: {self.stream_id}")
async def disconnect(self) -> None:
"""Closes the connection to the server.
@@ -93,5 +78,4 @@ class KafkaBroker(Broker):
iterator: AsyncIterator
An async iterator for the messages.
"""
-
return KafkaMessageFetcher(self.kafka_consumer)
diff --git
a/streampipes-client-python/streampipes/functions/broker/kafka_message_fetcher.py
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_message_fetcher.py
similarity index 94%
rename from
streampipes-client-python/streampipes/functions/broker/kafka_message_fetcher.py
rename to
streampipes-client-python/streampipes/functions/broker/kafka/kafka_message_fetcher.py
index d8e9fc3d2..efa7f31bd 100644
---
a/streampipes-client-python/streampipes/functions/broker/kafka_message_fetcher.py
+++
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_message_fetcher.py
@@ -46,5 +46,7 @@ class KafkaMessageFetcher:
return self
async def __anext__(self):
- msg = self.consumer.poll(0.1)
+ msg = None
+ while not msg:
+ msg = self.consumer.poll(0.1)
return KafkaMessage(msg.value())
diff --git
a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_publisher.py
similarity index 57%
copy from streampipes-client-python/streampipes/functions/broker/nats_broker.py
copy to
streampipes-client-python/streampipes/functions/broker/kafka/kafka_publisher.py
index 238e14bfd..f55337b5a 100644
--- a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
+++
b/streampipes-client-python/streampipes/functions/broker/kafka/kafka_publisher.py
@@ -14,20 +14,21 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+
import json
import logging
-from typing import Any, AsyncIterator, Dict
+from typing import Any, Dict
-from nats import connect
-from streampipes.functions.broker.broker import Broker
+from confluent_kafka import Producer # type: ignore
+from streampipes.functions.broker import Publisher
logger = logging.getLogger(__name__)
-class NatsBroker(Broker):
- """Implementation of the NatsBroker"""
+class KafkaPublisher(Publisher):
+ """Implementation of a publisher for Kafka"""
- async def _makeConnection(self, hostname: str, port: int) -> None:
+ async def _make_connection(self, hostname: str, port: int) -> None:
"""Helper function to connect to a server.
Parameters
@@ -42,22 +43,9 @@ class NatsBroker(Broker):
Returns
-------
None
-
"""
- self.nats_client = await connect(f"nats://{hostname}:{port}")
- if self.nats_client.connected_url is not None:
- logger.info(f"Connected to NATS at
{self.nats_client.connected_url.netloc}")
-
- async def createSubscription(self) -> None:
- """Creates a subscription to a data stream.
-
- Returns
- -------
- None
-
- """
- self.subscription = await self.nats_client.subscribe(self.topic_name)
- logger.info(f"Subscribed to stream: {self.stream_id}")
+ self.kafka_producer = Producer({"bootstrap.servers":
f"{hostname}:{port}"})
+ logger.info(f"Connecting to Kafka at {hostname}:{port}")
async def publish_event(self, event: Dict[str, Any]):
"""Publish an event to a connected data stream.
@@ -70,9 +58,9 @@ class NatsBroker(Broker):
Returns
-------
None
-
"""
- await self.nats_client.publish(subject=self.topic_name,
payload=json.dumps(event).encode("utf-8"))
+ self.kafka_producer.produce(topic=self.topic_name,
value=json.dumps(event).encode("utf-8"))
+ self.kafka_producer.flush()
async def disconnect(self) -> None:
"""Closes the connection to the server.
@@ -81,15 +69,4 @@ class NatsBroker(Broker):
-------
None
"""
- await self.nats_client.close()
logger.info(f"Stopped connection to stream: {self.stream_id}")
-
- def get_message(self) -> AsyncIterator:
- """Get the published messages of the subscription.
-
- Returns
- -------
- message_iterator: AsyncIterator
- An async iterator for the messages.
- """
- return self.subscription.messages
diff --git a/streampipes-client-python/streampipes/functions/broker/__init__.py
b/streampipes-client-python/streampipes/functions/broker/nats/__init__.py
similarity index 73%
copy from streampipes-client-python/streampipes/functions/broker/__init__.py
copy to streampipes-client-python/streampipes/functions/broker/nats/__init__.py
index 32d35fba3..cce3acad3 100644
--- a/streampipes-client-python/streampipes/functions/broker/__init__.py
+++ b/streampipes-client-python/streampipes/functions/broker/nats/__init__.py
@@ -14,16 +14,3 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from .broker import Broker
-from .kafka_broker import KafkaBroker
-from .nats_broker import NatsBroker
-
-from .broker_handler import SupportedBroker, get_broker # isort: skip
-
-__all__ = [
- "Broker",
- "KafkaBroker",
- "NatsBroker",
- "SupportedBroker",
- "get_broker",
-]
diff --git
a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
b/streampipes-client-python/streampipes/functions/broker/nats/nats_consumer.py
similarity index 70%
copy from streampipes-client-python/streampipes/functions/broker/nats_broker.py
copy to
streampipes-client-python/streampipes/functions/broker/nats/nats_consumer.py
index 238e14bfd..b57f0f35b 100644
--- a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
+++
b/streampipes-client-python/streampipes/functions/broker/nats/nats_consumer.py
@@ -14,20 +14,19 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-import json
import logging
-from typing import Any, AsyncIterator, Dict
+from typing import AsyncIterator
from nats import connect
-from streampipes.functions.broker.broker import Broker
+from streampipes.functions.broker import Consumer
logger = logging.getLogger(__name__)
-class NatsBroker(Broker):
- """Implementation of the NatsBroker"""
+class NatsConsumer(Consumer):
+ """Implementation of a consumer for NATS"""
- async def _makeConnection(self, hostname: str, port: int) -> None:
+ async def _make_connection(self, hostname: str, port: int) -> None:
"""Helper function to connect to a server.
Parameters
@@ -45,10 +44,9 @@ class NatsBroker(Broker):
"""
self.nats_client = await connect(f"nats://{hostname}:{port}")
- if self.nats_client.connected_url is not None:
- logger.info(f"Connected to NATS at
{self.nats_client.connected_url.netloc}")
+ logger.info(f"Connecting to NATS at {hostname}:{port}")
- async def createSubscription(self) -> None:
+ async def _create_subscription(self) -> None:
"""Creates a subscription to a data stream.
Returns
@@ -59,21 +57,6 @@ class NatsBroker(Broker):
self.subscription = await self.nats_client.subscribe(self.topic_name)
logger.info(f"Subscribed to stream: {self.stream_id}")
- async def publish_event(self, event: Dict[str, Any]):
- """Publish an event to a connected data stream.
-
- Parameters
- ----------
- event: Dict[str, Any]
- The event to be published.
-
- Returns
- -------
- None
-
- """
- await self.nats_client.publish(subject=self.topic_name,
payload=json.dumps(event).encode("utf-8"))
-
async def disconnect(self) -> None:
"""Closes the connection to the server.
diff --git
a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
b/streampipes-client-python/streampipes/functions/broker/nats/nats_publisher.py
similarity index 66%
rename from
streampipes-client-python/streampipes/functions/broker/nats_broker.py
rename to
streampipes-client-python/streampipes/functions/broker/nats/nats_publisher.py
index 238e14bfd..37ffa4e78 100644
--- a/streampipes-client-python/streampipes/functions/broker/nats_broker.py
+++
b/streampipes-client-python/streampipes/functions/broker/nats/nats_publisher.py
@@ -16,18 +16,18 @@
#
import json
import logging
-from typing import Any, AsyncIterator, Dict
+from typing import Any, Dict
from nats import connect
-from streampipes.functions.broker.broker import Broker
+from streampipes.functions.broker import Publisher
logger = logging.getLogger(__name__)
-class NatsBroker(Broker):
- """Implementation of the NatsBroker"""
+class NatsPublisher(Publisher):
+ """Implementation of a publisher for NATS"""
- async def _makeConnection(self, hostname: str, port: int) -> None:
+ async def _make_connection(self, hostname: str, port: int) -> None:
"""Helper function to connect to a server.
Parameters
@@ -45,19 +45,7 @@ class NatsBroker(Broker):
"""
self.nats_client = await connect(f"nats://{hostname}:{port}")
- if self.nats_client.connected_url is not None:
- logger.info(f"Connected to NATS at
{self.nats_client.connected_url.netloc}")
-
- async def createSubscription(self) -> None:
- """Creates a subscription to a data stream.
-
- Returns
- -------
- None
-
- """
- self.subscription = await self.nats_client.subscribe(self.topic_name)
- logger.info(f"Subscribed to stream: {self.stream_id}")
+ logger.info(f"Connecting to NATS at {hostname}:{port}")
async def publish_event(self, event: Dict[str, Any]):
"""Publish an event to a connected data stream.
@@ -83,13 +71,3 @@ class NatsBroker(Broker):
"""
await self.nats_client.close()
logger.info(f"Stopped connection to stream: {self.stream_id}")
-
- def get_message(self) -> AsyncIterator:
- """Get the published messages of the subscription.
-
- Returns
- -------
- message_iterator: AsyncIterator
- An async iterator for the messages.
- """
- return self.subscription.messages
diff --git
a/streampipes-client-python/streampipes/functions/broker/output_collector.py
b/streampipes-client-python/streampipes/functions/broker/output_collector.py
index fcae7d55e..a819539aa 100644
--- a/streampipes-client-python/streampipes/functions/broker/output_collector.py
+++ b/streampipes-client-python/streampipes/functions/broker/output_collector.py
@@ -17,7 +17,7 @@
import asyncio
from typing import Any, Coroutine, Dict
-from streampipes.functions.broker import get_broker
+from streampipes.functions.broker import Publisher, get_broker
from streampipes.model.resource.data_stream import DataStream
@@ -32,14 +32,14 @@ class OutputCollector:
Attributes
----------
- broker: Broker
- The broker instance that sends the data to StreamPipes
+ publisher: Publisher
+ The publisher instance that sends the data to StreamPipes
"""
def __init__(self, data_stream: DataStream) -> None:
- self.broker = get_broker(data_stream)
- self._run_coroutine(self.broker.connect(data_stream))
+ self.publisher: Publisher = get_broker(data_stream, is_publisher=True)
# type: ignore
+ self._run_coroutine(self.publisher.connect(data_stream))
def collect(self, event: Dict[str, Any]) -> None:
"""Publishes an event to the output stream.
@@ -53,7 +53,7 @@ class OutputCollector:
-------
None
"""
- self._run_coroutine(self.broker.publish_event(event))
+ self._run_coroutine(self.publisher.publish_event(event))
def disconnect(self) -> None:
"""Disconnects the broker of the output collector.
@@ -62,7 +62,7 @@ class OutputCollector:
-------
None
"""
- self._run_coroutine(self.broker.disconnect())
+ self._run_coroutine(self.publisher.disconnect())
@staticmethod
def _run_coroutine(coroutine: Coroutine) -> None:
diff --git a/streampipes-client-python/streampipes/functions/broker/__init__.py
b/streampipes-client-python/streampipes/functions/broker/publisher.py
similarity index 55%
copy from streampipes-client-python/streampipes/functions/broker/__init__.py
copy to streampipes-client-python/streampipes/functions/broker/publisher.py
index 32d35fba3..bdcb0d8e8 100644
--- a/streampipes-client-python/streampipes/functions/broker/__init__.py
+++ b/streampipes-client-python/streampipes/functions/broker/publisher.py
@@ -14,16 +14,29 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from .broker import Broker
-from .kafka_broker import KafkaBroker
-from .nats_broker import NatsBroker
+from abc import abstractmethod
+from typing import Any, Dict
-from .broker_handler import SupportedBroker, get_broker # isort: skip
+from streampipes.functions.broker import Broker
-__all__ = [
- "Broker",
- "KafkaBroker",
- "NatsBroker",
- "SupportedBroker",
- "get_broker",
-]
+
+class Publisher(Broker):
+ """Abstract implementation of a publisher for a broker.
+
+ A publisher allows to publish events to a data stream.
+ """
+
+ @abstractmethod
+ async def publish_event(self, event: Dict[str, Any]) -> None:
+ """Publish an event to a connected data stream.
+
+ Parameters
+ ----------
+ event: Dict[str, Any]
+ The event to be published.
+
+ Returns
+ -------
+ None
+ """
+ raise NotImplementedError # pragma: no cover
diff --git
a/streampipes-client-python/streampipes/functions/function_handler.py
b/streampipes-client-python/streampipes/functions/function_handler.py
index 83b88669b..deb2fc081 100644
--- a/streampipes-client-python/streampipes/functions/function_handler.py
+++ b/streampipes-client-python/streampipes/functions/function_handler.py
@@ -20,7 +20,7 @@ import logging
from typing import AsyncIterator, Dict, List
from streampipes.client.client import StreamPipesClient
-from streampipes.functions.broker import Broker, get_broker
+from streampipes.functions.broker import Broker, Consumer, get_broker
from streampipes.functions.registration import Registration
from streampipes.functions.utils.async_iter_handler import AsyncIterHandler
from streampipes.functions.utils.data_stream_context import DataStreamContext
@@ -77,7 +77,7 @@ class FunctionHandler:
# Get the data stream schema from the API
data_stream: DataStream =
self.client.dataStreamApi.get(stream_id) # type: ignore
# Get the broker
- broker = get_broker(data_stream)
+ broker: Consumer = get_broker(data_stream) # type: ignore
# Assign the functions, broker and schema to every stream
if stream_id in self.stream_contexts.keys():
self.stream_contexts[stream_id].add_function(streampipes_function)
@@ -110,7 +110,6 @@ class FunctionHandler:
broker = self.stream_contexts[stream_id].broker
# Connect the broker
await broker.connect(data_stream)
- await broker.createSubscription()
self.brokers.append(broker)
# Get the messages
messages[stream_id] = broker.get_message()
diff --git
a/streampipes-client-python/streampipes/functions/utils/async_iter_handler.py
b/streampipes-client-python/streampipes/functions/utils/async_iter_handler.py
index 2b69dfd6d..62a742d5b 100644
---
a/streampipes-client-python/streampipes/functions/utils/async_iter_handler.py
+++
b/streampipes-client-python/streampipes/functions/utils/async_iter_handler.py
@@ -39,7 +39,7 @@ class AsyncIterHandler:
"""
try:
return stream_id, await message.__anext__()
- except StopAsyncIteration:
+ except (StopAsyncIteration, RuntimeError):
return "stop", None
@staticmethod
diff --git
a/streampipes-client-python/streampipes/functions/utils/data_stream_context.py
b/streampipes-client-python/streampipes/functions/utils/data_stream_context.py
index 2132b5071..315a13cd0 100644
---
a/streampipes-client-python/streampipes/functions/utils/data_stream_context.py
+++
b/streampipes-client-python/streampipes/functions/utils/data_stream_context.py
@@ -16,7 +16,7 @@
#
from typing import List
-from streampipes.functions.broker.broker import Broker
+from streampipes.functions.broker import Consumer
from streampipes.functions.streampipes_function import StreamPipesFunction
from streampipes.model.resource.data_stream import DataStream
@@ -30,11 +30,11 @@ class DataStreamContext:
StreamPipes Functions which require the data of this data stream.
schema: DataStream
The schema of this data stream.
- broker: Broker
- The broker to connect to this data stream.
+ broker: Consumer
+ The consumer to connect to this data stream.
"""
- def __init__(self, functions: List[StreamPipesFunction], schema:
DataStream, broker: Broker) -> None:
+ def __init__(self, functions: List[StreamPipesFunction], schema:
DataStream, broker: Consumer) -> None:
self.functions = functions
self.schema = schema
self.broker = broker
diff --git
a/streampipes-client-python/streampipes/functions/utils/data_stream_generator.py
b/streampipes-client-python/streampipes/functions/utils/data_stream_generator.py
index e43d15cc4..ee2ad3618 100644
---
a/streampipes-client-python/streampipes/functions/utils/data_stream_generator.py
+++
b/streampipes-client-python/streampipes/functions/utils/data_stream_generator.py
@@ -18,7 +18,13 @@
from enum import Enum
from typing import Dict, Optional
-from streampipes.model.common import EventProperty, EventSchema
+from streampipes.functions.broker import SupportedBroker
+from streampipes.model.common import (
+ EventGrounding,
+ EventProperty,
+ EventSchema,
+ TransportProtocol,
+)
from streampipes.model.resource.data_stream import DataStream
@@ -44,7 +50,12 @@ class RuntimeType(Enum):
# TODO Use an more general approach to create a data stream
-def create_data_stream(name: str, attributes: Dict[str, str], stream_id:
Optional[str] = None):
+def create_data_stream(
+ name: str,
+ attributes: Dict[str, str],
+ stream_id: Optional[str] = None,
+ broker: SupportedBroker = SupportedBroker.NATS,
+):
"""Creates a data stream
Parameters
@@ -83,6 +94,19 @@ def create_data_stream(name: str, attributes: Dict[str,
str], stream_id: Optiona
]
)
- if not stream_id:
- return DataStream(name=name, event_schema=event_schema)
- return DataStream(element_id=stream_id, name=name,
event_schema=event_schema)
+ transport_protocols = [TransportProtocol()]
+ if broker == SupportedBroker.KAFKA:
+ transport_protocols = [
+ TransportProtocol(
+
class_name="org.apache.streampipes.model.grounding.KafkaTransportProtocol", #
type: ignore
+ broker_hostname="kafka",
+ port=9092,
+ )
+ ]
+
+ data_stream = DataStream(
+ name=name, event_schema=event_schema,
event_grounding=EventGrounding(transport_protocols=transport_protocols)
+ )
+ if stream_id:
+ data_stream.element_id = stream_id
+ return data_stream
diff --git a/streampipes-client-python/tests/client/test_endpoint.py
b/streampipes-client-python/tests/client/test_endpoint.py
index 16a2b6d9d..5e4c0e1f9 100644
--- a/streampipes-client-python/tests/client/test_endpoint.py
+++ b/streampipes-client-python/tests/client/test_endpoint.py
@@ -28,7 +28,7 @@ from streampipes.client.config import StreamPipesClientConfig
from streampipes.client.credential_provider import StreamPipesApiKeyCredentials
from streampipes.endpoint.endpoint import MessagingEndpoint,
_error_code_to_message
from streampipes.endpoint.exceptions import MessagingEndpointNotConfiguredError
-from streampipes.functions.broker.nats_broker import NatsBroker
+from streampipes.functions.broker import NatsConsumer
from streampipes.model.container.resource_container import (
StreamPipesDataModelError,
StreamPipesResourceContainerJSONError,
@@ -421,9 +421,9 @@ class TestMessagingEndpoint(TestCase):
def test_messaging_endpoint_happy_path(self):
demo_endpoint = MessagingEndpoint(parent_client=self.client)
- demo_endpoint.configure(broker=NatsBroker())
+ demo_endpoint.configure(broker=NatsConsumer())
- self.assertTrue(isinstance(demo_endpoint.broker, NatsBroker))
+ self.assertTrue(isinstance(demo_endpoint.broker, NatsConsumer))
def test_messaging_endpoint_missing_configure(self):
demo_endpoint = MessagingEndpoint(parent_client=self.client)
diff --git a/streampipes-client-python/tests/functions/test_function_handler.py
b/streampipes-client-python/tests/functions/test_function_handler.py
index 6b544f656..8ff325299 100644
--- a/streampipes-client-python/tests/functions/test_function_handler.py
+++ b/streampipes-client-python/tests/functions/test_function_handler.py
@@ -14,13 +14,18 @@
# 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, patch
+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
@@ -83,7 +88,7 @@ class TestFunctionOutput(StreamPipesFunction):
self.stopped = True
-class TestMessage:
+class TestNatsMessage:
def __init__(self, data) -> None:
self.data = JSONEncoder().encode(data).encode()
@@ -99,115 +104,53 @@ class TestMessageIterator:
async def __anext__(self):
if self.i < len(self.test_data) - 1:
self.i += 1
- return TestMessage(self.test_data[self.i])
+ 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:
- # set example responses from endpoints
- self.data_stream: Dict[str, Any] = {
- "elementId": "urn:streampipes.apache.org:eventstream:uPDKLI",
- "name": "Test",
- "description": "",
- "iconUrl": None,
- "appId": None,
- "includesAssets": False,
- "includesLocales": False,
- "includedAssets": [],
- "includedLocales": [],
- "applicationLinks": [],
- "internallyManaged": True,
- "connectedTo": None,
- "eventGrounding": {
- "elementId":
"urn:streampipes.apache.org:spi:eventgrounding:TwGIQA",
- "transportProtocols": [
- {
- "@class":
"org.apache.streampipes.model.grounding.NatsTransportProtocol",
- "elementId":
"urn:streampipes.apache.org:spi:natstransportprotocol:VJkHmZ",
- "brokerHostname": "nats",
- "topicDefinition": {
- "elementId":
"urn:streampipes.apache.org:spi:simpletopicdefinition:QzCiFI",
- "actualTopicName": "test1",
- },
- "port": 4222,
- }
- ],
- "transportFormats": [
- {
- "elementId":
"urn:streampipes.apache.org:spi:transportformat:CMGsLP",
- "rdfType":
["http://sepa.event-processing.org/sepa#json"],
- }
- ],
- },
- "eventSchema": {
- "elementId":
"urn:streampipes.apache.org:spi:eventschema:rARlLX",
- "eventProperties": [
- {
- "elementId":
"urn:streampipes.apache.org:spi:eventpropertyprimitive:yogPNV",
- "label": "Density",
- "description": "Denotes the current density of the
fluid",
- "runtimeName": "density",
- "required": False,
- "domainProperties": ["http://schema.org/Number"],
- "eventPropertyQualities": [],
- "requiresEventPropertyQualities": [],
- "propertyScope": "MEASUREMENT_PROPERTY",
- "index": 5,
- "runtimeId": None,
- "runtimeType":
"http://www.w3.org/2001/XMLSchema#float",
- "measurementUnit": None,
- "valueSpecification": None,
- },
- {
- "elementId":
"urn:streampipes.apache.org:spi:eventpropertyprimitive:GjZgFg",
- "label": "Temperature",
- "description": "Denotes the current temperature in
degrees celsius",
- "runtimeName": "temperature",
- "required": False,
- "domainProperties": ["http://schema.org/Number"],
- "eventPropertyQualities": [],
- "requiresEventPropertyQualities": [],
- "propertyScope": "MEASUREMENT_PROPERTY",
- "index": 4,
- "runtimeId": None,
- "runtimeType":
"http://www.w3.org/2001/XMLSchema#float",
- "measurementUnit":
"http://codes.wmo.int/common/unit/degC",
- "valueSpecification": {
- "elementId":
"urn:streampipes.apache.org:spi:quantitativevalue:ZQSJfk",
- "minValue": 0,
- "maxValue": 100,
- "step": 0.1,
- },
- },
- {
- "@class":
"org.apache.streampipes.model.schema.EventPropertyPrimitive",
- "elementId":
"urn:streampipes.apache.org:spi:eventpropertyprimitive:NRKdap",
- "label": "Timestamp",
- "description": "The current timestamp value",
- "runtimeName": "timestamp",
- "required": False,
- "domainProperties": ["http://schema.org/DateTime"],
- "eventPropertyQualities": [],
- "requiresEventPropertyQualities": [],
- "propertyScope": "HEADER_PROPERTY",
- "index": 0,
- "runtimeId": None,
- "runtimeType": "http://www.w3.org/2001/XMLSchema#long",
- "measurementUnit": None,
- "valueSpecification": None,
- },
- ],
- },
- "measurementCapability": None,
- "measurementObject": None,
- "index": 0,
- "correspondingAdapterId":
"urn:streampipes.apache.org:spi:org.apache.streampipes.connect.iiot...",
- "category": None,
- "uri": "urn:streampipes.apache.org:eventstream:uPDKLI",
- "dom": 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},
@@ -228,17 +171,49 @@ class TestFunctionHandler(TestCase):
{"density": 3.6, "temperature": 30.4, "timestamp": 1670000007000},
]
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.disconnect",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.createSubscription",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker._makeConnection",
autospec=True)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.get_message",
autospec=True)
+ @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)
+ def test_function_handler_nats(self, 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
+
+ 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()
+ 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)
- def test_function_handler(self, http_session: MagicMock, nats_broker:
MagicMock, *args: Tuple[AsyncMock]):
+ def test_function_handler_kafka(self, http_session: MagicMock, connection:
MagicMock):
http_session_mock = MagicMock()
- http_session_mock.get.return_value.json.return_value = self.data_stream
+ http_session_mock.get.return_value.json.return_value =
self.data_stream_kafka
http_session.return_value = http_session_mock
- nats_broker.return_value = TestMessageIterator(self.test_stream_data1)
+ 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(
@@ -255,7 +230,7 @@ class TestFunctionHandler(TestCase):
self.assertEqual(test_function.context.client, client)
self.assertDictEqual(
- test_function.context.schema, {self.data_stream["elementId"]:
DataStream(**self.data_stream)}
+ 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)
@@ -263,17 +238,36 @@ class TestFunctionHandler(TestCase):
self.assertListEqual(test_function.data, self.test_stream_data1)
self.assertTrue(test_function.stopped)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.disconnect",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.createSubscription",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker._makeConnection",
autospec=True)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.get_message",
autospec=True)
+ @patch("streampipes.functions.broker.nats.nats_consumer.connect",
autospec=True)
+ @patch("streampipes.client.client.Session", autospec=True)
+ def test_function_handler_unsupported_broker(self, 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
+
+ client = StreamPipesClient(
+ client_config=StreamPipesClientConfig(
+
credential_provider=StreamPipesApiKeyCredentials(username="user",
api_key="key"),
+ host_address="localhost",
+ )
+ )
+
+ registration = Registration()
+ test_function = TestFunction()
+ 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.endpoint.APIEndpoint.get", autospec=True)
- def test_function_handler_two_streams(self, endpoint: MagicMock,
nats_broker: MagicMock, *args: Tuple[AsyncMock]):
+ def test_two_streams_nats(self, 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)
+ return DataStream(**self.data_stream_nats)
elif stream_id == "urn:streampipes.apache.org:eventstream:HHoidJ":
- data_stream = DataStream(**self.data_stream)
+ 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
@@ -318,14 +312,13 @@ class TestFunctionHandler(TestCase):
self.assertListEqual(test_function2.data2, self.test_stream_data2)
self.assertTrue(test_function2.stopped)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.disconnect",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.createSubscription",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker._makeConnection",
autospec=True)
+ @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.nats_broker.NatsBroker.get_message",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.publish_event",
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)
- def test_function_output_stream(
+ def test_function_output_stream_nats(
self,
http_session: MagicMock,
pulish_event: MagicMock,
@@ -334,7 +327,7 @@ class TestFunctionHandler(TestCase):
*args: Tuple[AsyncMock]
):
http_session_mock = MagicMock()
- http_session_mock.get.return_value.json.return_value = self.data_stream
+ http_session_mock.get.return_value.json.return_value =
self.data_stream_nats
http_session.return_value = http_session_mock
output_events = []
@@ -364,7 +357,66 @@ class TestFunctionHandler(TestCase):
self.assertEqual(test_function.context.client, client)
self.assertDictEqual(
- test_function.context.schema, {self.data_stream["elementId"]:
DataStream(**self.data_stream)}
+ 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)
+ def test_function_output_stream_kafka(
+ self,
+ 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
+
+ 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", attributes={"number": RuntimeType.INTEGER.value},
broker=SupportedBroker.KAFKA
+ )
+ test_function = TestFunctionOutput(
+
function_definition=FunctionDefinition().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)
diff --git a/streampipes-client-python/tests/functions/test_river_function.py
b/streampipes-client-python/tests/functions/test_river_function.py
index a7035d003..a850501ef 100644
--- a/streampipes-client-python/tests/functions/test_river_function.py
+++ b/streampipes-client-python/tests/functions/test_river_function.py
@@ -67,12 +67,14 @@ class TestRiverFunction(TestCase):
{"number": 10.6, "bool": False, "timestamp": 1670000005000},
]
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.disconnect",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.createSubscription",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker._makeConnection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsPublisher.disconnect",
autospec=True)
+ @patch("streampipes.functions.broker.NatsPublisher._make_connection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer.disconnect",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer._make_connection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer._create_subscription",
autospec=True)
@patch("streampipes.functions.streampipes_function.time", autospec=True)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.get_message",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.publish_event",
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)
def test_river_function_unsupervised(
self,
@@ -130,12 +132,14 @@ class TestRiverFunction(TestCase):
self.assertListEqual(model.data_x, self.test_stream_data)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.disconnect",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.createSubscription",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker._makeConnection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsPublisher.disconnect",
autospec=True)
+ @patch("streampipes.functions.broker.NatsPublisher._make_connection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer.disconnect",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer._make_connection",
autospec=True)
+ @patch("streampipes.functions.broker.NatsConsumer._create_subscription",
autospec=True)
@patch("streampipes.functions.streampipes_function.time", autospec=True)
- @patch("streampipes.functions.broker.nats_broker.NatsBroker.get_message",
autospec=True)
-
@patch("streampipes.functions.broker.nats_broker.NatsBroker.publish_event",
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)
def test_river_function_supervised(
self,