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,

Reply via email to