fix: add commit and commited functions to consumer test client. Closes …

…#61 (#63)

Co-authored-by: Marcos Schroh <[email protected]>

docs/test_client.md

@@ -115,6 +115,48 @@ async def test_add_event_on_consume():
     The `TestStreamClient.send` coroutine is used instead.
     This allows to test `streams` without having producer code in your application
 
+### Testing the Commit
+
+In some cases your stream will commit, in this situation checking the commited partitions can be useful.
+
+```python
+import pytest
+from kstreams.test_utils import TestStreamClient
+
+from .example import produce, stream_engine
+
+topic_name = "local--kstreams-marcos"
+value = b'{"message": "Hello world!"}'
+name = "my-stream"
+key = "1"
+partition = 2
+tp = structs.TopicPartition(
+    topic=topic_name,
+    partition=partition,
+)
+total_events = 10
+
+@stream_engine.stream(topic_name, name=name)
+async def my_stream(stream: Stream):
+    async for cr in stream:
+        # commit every time that an event arrives
+        await stream.consumer.commit({tp: cr.offset})
+
+
+# test the code
+client = TestStreamClient(stream_engine)
+
+@pytest.mark.asyncio
+async def test_consumer_commit(stream_engine: StreamEngine):
+    async with client:
+        for _ in range(0, total_events):
+            await client.send(topic_name, partition=partition, value=value, key=key)
+
+    # check that everything was commited
+    stream = stream_engine.get_stream(name)
+    assert (await stream.consumer.committed(tp)) == total_events
+```
+
 ### E2E test
 
 In the previous code example the application produces to and consumes from the same topic, then `TestStreamClient.send` is not needed because the `engine.send` is producing. For those situation you can just use your `producer` code and check that certain code was called.

kstreams/engine.py

@@ -46,7 +46,7 @@ async def send(
         topic: str,
         value: Any = None,
         key: Any = None,
-        partition: Optional[str] = None,
+        partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         headers: Optional[Headers] = None,
         serializer: Optional[Serializer] = None,

kstreams/streams.py

@@ -73,6 +73,9 @@ async def stop(self) -> None:
             self._consumer_task.cancel()
 
     async def start(self) -> Optional[AsyncGenerator]:
+        if self.running:
+            return None
+
         async def func_wrapper(func):
             try:
                 # await for the end user coroutine

kstreams/test_utils/test_clients.py

@@ -65,6 +65,7 @@ def __init__(self, *topics: str, group_id: Optional[str] = None, **kwargs) -> No
         self.topics: Tuple[str, ...] = topics
         self._group_id: Optional[str] = group_id
         self._assigments: List[TopicPartition] = []
+        self.partitions_committed: Dict[TopicPartition, int] = {}
 
         for topic_name in topics:
             TopicManager.create(topic_name, consumer=self)
@@ -90,6 +91,15 @@ async def position(self, topic_partition: TopicPartition) -> int:
     def highwater(self, topic_partition: TopicPartition) -> int:
         return self.last_stable_offset(topic_partition)
 
+    async def commit(self, offsets: Optional[Dict[TopicPartition, int]] = None) -> None:
+        if offsets is not None:
+            for topic_partition, offset in offsets.items():
+                self.partitions_committed[topic_partition] = offset
+        return None
+
+    async def committed(self, topic_partition: TopicPartition) -> Optional[int]:
+        return self.partitions_committed.get(topic_partition)
+
     async def getone(
         self,
     ) -> Optional[ConsumerRecord]:  # The return type must be fixed later on

kstreams/test_utils/test_utils.py

@@ -70,7 +70,7 @@ async def send(
         topic: str,
         value: Any = None,
         key: Optional[Any] = None,
-        partition: Optional[str] = None,
+        partition: Optional[int] = None,
         timestamp_ms: Optional[int] = None,
         headers: Optional[Headers] = None,
         serializer: Optional[Serializer] = None,

tests/test_client.py

@@ -3,6 +3,7 @@
 import pytest
 
 from kstreams import StreamEngine
+from kstreams.streams import Stream
 from kstreams.test_utils import (
     TestConsumer,
     TestProducer,
@@ -87,6 +88,34 @@ async def test_topic_created(stream_engine: StreamEngine):
     assert consumer_record.key == key
 
 
+@pytest.mark.asyncio
+async def test_consumer_commit(stream_engine: StreamEngine):
+    topic_name = "local--kstreams-marcos"
+    value = b'{"message": "Hello world!"}'
+    name = "my-stream"
+    key = "1"
+    partition = 2
+    tp = structs.TopicPartition(
+        topic=topic_name,
+        partition=partition,
+    )
+    total_events = 10
+
+    @stream_engine.stream(topic_name, name=name)
+    async def my_stream(stream: Stream):
+        async for cr in stream:
+            await stream.consumer.commit({tp: cr.offset})
+
+    client = TestStreamClient(stream_engine)
+    async with client:
+        for _ in range(0, total_events):
+            await client.send(topic_name, partition=partition, value=value, key=key)
+
+    # check that everything was commited
+    stream = stream_engine.get_stream(name)
+    assert (await stream.consumer.committed(tp)) == total_events
+
+
 @pytest.mark.asyncio
 async def test_e2e_example():
     """