sentry-arroyo 2.19.4__py3-none-any.whl → 2.19.6__py3-none-any.whl

arroyo/backends/kafka/commit.py

@@ -97,6 +97,6 @@ class CommitCodec(Codec[KafkaPayload, Commit]):
 
         if max_times_to_log_legacy_message > 0:
             max_times_to_log_legacy_message -= 1
-            logger.warn(f"Legacy commit message found: {commit}")
+            logger.warning("Legacy commit message found: %s", commit)
 
         return commit

arroyo/backends/kafka/consumer.py

@@ -161,6 +161,10 @@ class KafkaConsumer(Consumer[KafkaPayload]):
         )
 
         configuration = dict(configuration)
+        self.__is_incremental = (
+            configuration.get("partition.assignment.strategy") == "cooperative-sticky"
+            or configuration.get("group.protocol") == "consumer"
+        )
         auto_offset_reset = configuration.get("auto.offset.reset", "largest")
 
         # This is a special flag that controls the auto offset behavior for
@@ -269,6 +273,10 @@ class KafkaConsumer(Consumer[KafkaPayload]):
         def assignment_callback(
             consumer: ConfluentConsumer, partitions: Sequence[ConfluentTopicPartition]
         ) -> None:
+            if not partitions:
+                logger.info("skipping empty assignment")
+                return
+
             self.__state = KafkaConsumerState.ASSIGNING
 
             try:
@@ -451,12 +459,14 @@ class KafkaConsumer(Consumer[KafkaPayload]):
 
     def __assign(self, offsets: Mapping[Partition, int]) -> None:
         self.__validate_offsets(offsets)
-        self.__consumer.assign(
-            [
-                ConfluentTopicPartition(partition.topic.name, partition.index, offset)
-                for partition, offset in offsets.items()
-            ]
-        )
+        partitions = [
+            ConfluentTopicPartition(partition.topic.name, partition.index, offset)
+            for partition, offset in offsets.items()
+        ]
+        if self.__is_incremental:
+            self.__consumer.incremental_assign(partitions)
+        else:
+            self.__consumer.assign(partitions)
         self.__offsets.update(offsets)
 
     def seek(self, offsets: Mapping[Partition, int]) -> None:

arroyo/dlq.py

@@ -315,11 +315,11 @@ class BufferedMessages(Generic[TStrategyPayload]):
 
         return None
 
-    def reset(self) -> None:
+    def remove(self, partition: Partition) -> None:
         """
-        Reset the buffer.
+        Remove a revoked partition from the buffer.
         """
-        self.__buffered_messages = defaultdict(deque)
+        self.__buffered_messages.pop(partition, None)
 
 
 class DlqPolicyWrapper(Generic[TStrategyPayload]):
@@ -343,9 +343,9 @@ class DlqPolicyWrapper(Generic[TStrategyPayload]):
                 ]
             ],
         ] = defaultdict(deque)
-        self.reset_offsets({})
+        self.reset_dlq_limits({})
 
-    def reset_offsets(self, assignment: Mapping[Partition, int]) -> None:
+    def reset_dlq_limits(self, assignment: Mapping[Partition, int]) -> None:
         """
         Called on consumer assignment
         """

arroyo/processing/processor.py

@@ -238,16 +238,23 @@ class StreamProcessor(Generic[TStrategyPayload]):
                 "arroyo.consumer.partitions_assigned.count", len(partitions)
             )
 
-            self.__buffered_messages.reset()
+            current_partitions = dict(self.__consumer.tell())
+            current_partitions.update(partitions)
+
             if self.__dlq_policy:
-                self.__dlq_policy.reset_offsets(partitions)
-            if partitions:
+                self.__dlq_policy.reset_dlq_limits(current_partitions)
+            if current_partitions:
                 if self.__processing_strategy is not None:
-                    logger.exception(
+                    # TODO: for cooperative-sticky rebalancing this can happen
+                    # quite often. we should port the changes to
+                    # ProcessingStrategyFactory that we made in Rust: Remove
+                    # create_with_partitions, replace with create +
+                    # update_partitions
+                    logger.warning(
                         "Partition assignment while processing strategy active"
                     )
                     _close_strategy()
-                _create_strategy(partitions)
+                _create_strategy(current_partitions)
 
         @_rdkafka_callback(metrics=self.__metrics_buffer)
         def on_partitions_revoked(partitions: Sequence[Partition]) -> None:
@@ -278,6 +285,9 @@ class StreamProcessor(Generic[TStrategyPayload]):
                 except RuntimeError:
                     pass
 
+            for partition in partitions:
+                self.__buffered_messages.remove(partition)
+
             # Partition revocation can happen anytime during the consumer lifecycle and happen
             # multiple times. What we want to know is that the consumer is not stuck somewhere.
             # The presence of this message as the last message of a consumer

arroyo/utils/clock.py

@@ -28,7 +28,7 @@ class SystemClock(Clock):
         time.sleep(duration)
 
 
-class TestingClock(Clock):
+class MockedClock(Clock):
     """
     A clock implementation that uses a stable time for testing. To advance
     the time, use the ``sleep`` method.

{sentry_arroyo-2.19.4.dist-info → sentry_arroyo-2.19.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sentry-arroyo
-Version: 2.19.4
+Version: 2.19.6
 Summary: Arroyo is a Python library for working with streaming data.
 Home-page: https://github.com/getsentry/arroyo
 Author: Sentry

{sentry_arroyo-2.19.4.dist-info → sentry_arroyo-2.19.6.dist-info}/RECORD

@@ -1,22 +1,22 @@
 arroyo/__init__.py,sha256=fcpHZd2P3MxWl6PJJ8n__fM_NRIfiUE8tKN-orv6lb0,187
 arroyo/commit.py,sha256=oFihWUW8fLsjomWh0o085qIHe9vwVNgoOJC6JQdFM7M,2235
-arroyo/dlq.py,sha256=2_IlVfG_bizSSPOyGXLaI6yJCDoAWMkWx42ah0JTJr0,15260
+arroyo/dlq.py,sha256=LkWbgI6jQBaiGw8mPzfG-vZ62at5IBsFivRcCDRDNK0,15315
 arroyo/errors.py,sha256=IbtoIbz_m5QrxNRBLOxiy-hOfJQTEwNPCyq6yqedJYk,1059
 arroyo/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arroyo/types.py,sha256=sLY0x030np4UmbaW5C1KH1se7Z2pjQiPvAe5x2sXf7A,5684
 arroyo/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arroyo/backends/abstract.py,sha256=PicUWWsgH-41XpHDcuq86P4PpPx7GWgLWdwnw4B0obo,7582
 arroyo/backends/kafka/__init__.py,sha256=TZ0omd3LFXcZUaN_soFTuTgbjNEQYF1mF_i6_KIfCNo,306
-arroyo/backends/kafka/commit.py,sha256=JV1C8hbTUTWyDIU1qcaQkUzQdc1j-APXIjYUJJRpb1o,3056
+arroyo/backends/kafka/commit.py,sha256=LPsjvX5PPXR62DT6sa5GuSF78qk9F_L--Fz4kw7-m-s,3060
 arroyo/backends/kafka/configuration.py,sha256=D9zre9H2xagUsk7qBA-bm38V3_4Mg_X5hpKsKx2BkM8,3048
-arroyo/backends/kafka/consumer.py,sha256=LUUp-jffNfgW8Qerw8Yv8IhVabw-oONrkIFrjDVcauM,27304
+arroyo/backends/kafka/consumer.py,sha256=sWqmn8jDf-wfZOXRc7vrotEepfeRbtdpzoP46VDf49A,27721
 arroyo/backends/local/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arroyo/backends/local/backend.py,sha256=pQeZRUo9KDWlOT1pICWM0MvS39QuJqHnN5ryrmcxnSc,13878
 arroyo/backends/local/storages/__init__.py,sha256=AGYujdAAcn3osoj9jq84IzTywYbkIDv9wRg2rLhLXeg,104
 arroyo/backends/local/storages/abstract.py,sha256=1qVQp6roxHkK6XT2aklZyZk1qq7RzcPN6Db_CA5--kg,2901
 arroyo/backends/local/storages/memory.py,sha256=AoKDsVZzBXkOJyWArKWp3vfGfU9xLlKFXE9gsJiMIzQ,2613
 arroyo/processing/__init__.py,sha256=vZVg0wJvJfoVzlzGvnL59bT6YNIRJNQ5t7oU045Qbk4,87
-arroyo/processing/processor.py,sha256=kIb_IXoRyy8j85H1--KheQYxeJzBTTZ6V9yc59LkpYw,18475
+arroyo/processing/processor.py,sha256=rbguYCg51rX8PVmhxUbf-IXZZulX2pVJUVVHJgU76bw,18998
 arroyo/processing/strategies/__init__.py,sha256=EU_JMb54eOxMxaC5mIFpI-sAF-X2ZScbE8czBZ7bQkY,1106
 arroyo/processing/strategies/abstract.py,sha256=nu7juEz_aQmQIH35Z8u--FBuLjkK8_LQ1hIG2xpw9AA,4808
 arroyo/processing/strategies/batching.py,sha256=s89xC6lQpBseEaApu1iNTipXGKeO95OMwinj2VBKn9s,4778
@@ -33,7 +33,7 @@ arroyo/processing/strategies/run_task_in_threads.py,sha256=f1sb2AG-BLz11X78jfhtE
 arroyo/processing/strategies/run_task_with_multiprocessing.py,sha256=XhLRiMXFhZTRlxCTTmeI6SSTxy9SwiDHHiqq6mow-Z0,34409
 arroyo/processing/strategies/unfold.py,sha256=wZDNdMo1Ln27P1tUvLO2svhL-JDs2ZCyIFF2cq0DxVQ,3400
 arroyo/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arroyo/utils/clock.py,sha256=xZE-amURj5wiuQXRpdoakB_osmEVf3llNpMLV0lVX8g,945
+arroyo/utils/clock.py,sha256=r2EMO4nL5qIb1xnAd1sTAk2yK1UltyUi04lk5BqWKIc,944
 arroyo/utils/codecs.py,sha256=x-8SJK0GLTOH4c_k24K97JPjBckxyQJcSpgoEViGUy0,541
 arroyo/utils/concurrent.py,sha256=dbdPinjqmxCQ7izUGFNbGjB3OxfSIO01bnCSTANaVOE,1187
 arroyo/utils/logging.py,sha256=Y1PnhYcI9XNNEK0H13Ct2xKLr2Niuw0dxayc6sWnui8,606
@@ -47,12 +47,12 @@ examples/transform_and_produce/batched.py,sha256=st2R6qTneAtV0JFbKP30Ti3sJDYj8Jk
 examples/transform_and_produce/script.py,sha256=8kSMIjQNqGYEVyE0PvrfJh-a_UYCrJSstTp_De7kyyg,2306
 examples/transform_and_produce/simple.py,sha256=H7xqxItjl4tx34wVW5dy6mB9G39QucAtxkJSBzVmjgA,1637
 tests/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/backends/mixins.py,sha256=I8oDRWgUhQEo6P6v1rYqrwiqcaigxSaJIMSHd7bntj4,17228
+tests/backends/mixins.py,sha256=31t7B4a8cJ65XeUzv4ZoG9aPU1lTpDZXgkQrEyKhyqE,20131
 tests/backends/test_commit.py,sha256=iTHfK1qsBxim0XwxgMvNNSMqDUMEHoYkYBDcgxGBFbs,831
-tests/backends/test_kafka.py,sha256=BP1OBWmciG-4yWWqdcrAV3406azFeHPmnICqUSLD3YI,11151
-tests/backends/test_local.py,sha256=JbWvB5ssNp9fggHDQRT53O_qpuQAN0GgwTLFXWLzIT4,3365
+tests/backends/test_kafka.py,sha256=cB_GIY1yWT9w2VB1PWrvTr0lCa46r5uDfxisDFrcSmg,12089
+tests/backends/test_local.py,sha256=Mfd4DFuWVSVtl1GomQ6TIoWuJNcAliKqKU0BShPlEMY,3363
 tests/processing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/processing/test_processor.py,sha256=5KT9I7rcI5fIERVFdUZQJ8R8B7c4ndHOblfiTbylDlE,20428
+tests/processing/test_processor.py,sha256=PDCrmhWAt_wZKwYlBzDuEGsd8PnjKM-p4ySaPvVd11k,20781
 tests/processing/strategies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/processing/strategies/test_all.py,sha256=ahAF-nbdmqVkYGNCg0OFCD6fzNTA-XxYrW8NQHajCDU,10167
 tests/processing/strategies/test_batching.py,sha256=nyyX0y6qYHX7jT4gCgsUjT5RzBMDrBp790SCmOizQ0Q,11787
@@ -61,7 +61,7 @@ tests/processing/strategies/test_commit.py,sha256=ruGcBa2qBf0BoxdGZh2hE4Fa_XnOX7
 tests/processing/strategies/test_filter.py,sha256=hbuWDWF1Eie9ditJa6dZ3VgsJ-ODU7EADawD-gHXwa8,8411
 tests/processing/strategies/test_guard.py,sha256=fMx2wMlbVDDX4qqXuNRqzcLf5H4seTWw2kk4tBYK9d8,1361
 tests/processing/strategies/test_noop.py,sha256=a5noKxkWtY7SR_gcfu4oSNk2zvJA8Zmv7PncAdWNryk,512
-tests/processing/strategies/test_produce.py,sha256=YIExrxIhhNtekj3tr9KjSoX7EGFnxUQyhYrwQvlzq-k,2174
+tests/processing/strategies/test_produce.py,sha256=UQ03quIAvfnsg8Og7US6D4ERs-J8nCT12bHVnjHqxRw,2172
 tests/processing/strategies/test_reduce.py,sha256=crPFtGp7cyD8QOsmfVsyYh8KLOTzb8ryI7XtYg0vQSQ,1101
 tests/processing/strategies/test_run_task.py,sha256=bWIy4U6QyOBtqdiJdGLMAadlEME-W2aE_ZzDbU_BsGo,2805
 tests/processing/strategies/test_run_task_in_threads.py,sha256=5nwzF1iV6MTK1xETzWvMEOwAcZWrMOQaIPSWbiAjKFo,1457
@@ -70,9 +70,9 @@ tests/processing/strategies/test_unfold.py,sha256=Qic2Y2Un9EYBiW6E84YRvpyNctDNbf
 tests/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/utils/test_concurrent.py,sha256=Gwdzym2UZ1HO3rhOSGmzxImWcLFygY8P7MXHT3Q0xTE,455
 tests/utils/test_metrics.py,sha256=bI0EtGgPokMQyEqX58i0-8zvLfxRP2nWaWr2wLMaJ_o,917
-tests/utils/test_retries.py,sha256=qSXalR9TvjvbPpMAinhmJvpmoVo8Ybv07c5VgQZZeI8,3081
-sentry_arroyo-2.19.4.dist-info/LICENSE,sha256=0Ng3MFdEcnz0sVD1XvGBBzbavvNp_7OAM5yVObB46jU,10829
-sentry_arroyo-2.19.4.dist-info/METADATA,sha256=jDFHiPwdPjQtmegxMa78Xgq7QiMVASqgmkIE3S9vpdU,1989
-sentry_arroyo-2.19.4.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-sentry_arroyo-2.19.4.dist-info/top_level.txt,sha256=DVdMZKysL_iIxm5aY0sYgZtP5ZXMg9YBaBmGQHVmDXA,22
-sentry_arroyo-2.19.4.dist-info/RECORD,,
+tests/utils/test_retries.py,sha256=AxJLkXWeL9AjHv_p1n0pe8CXXJp24ZQIuYBHfNcmiz4,3075
+sentry_arroyo-2.19.6.dist-info/LICENSE,sha256=0Ng3MFdEcnz0sVD1XvGBBzbavvNp_7OAM5yVObB46jU,10829
+sentry_arroyo-2.19.6.dist-info/METADATA,sha256=smHYlB9nbo3gVdL5EodtxEihujwnzo4GYsl8VGao52Y,1989
+sentry_arroyo-2.19.6.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+sentry_arroyo-2.19.6.dist-info/top_level.txt,sha256=DVdMZKysL_iIxm5aY0sYgZtP5ZXMg9YBaBmGQHVmDXA,22
+sentry_arroyo-2.19.6.dist-info/RECORD,,

tests/backends/mixins.py

@@ -2,7 +2,7 @@ import time
 import uuid
 from abc import ABC, abstractmethod
 from contextlib import closing
-from typing import ContextManager, Generic, Iterator, Mapping, Optional, Sequence
+from typing import Any, ContextManager, Generic, Iterator, Mapping, Optional, Sequence
 from unittest import mock
 
 import pytest
@@ -14,6 +14,8 @@ from tests.assertions import assert_changes, assert_does_not_change
 
 
 class StreamsTestMixin(ABC, Generic[TStrategyPayload]):
+    cooperative_sticky = False
+
     @abstractmethod
     def get_topic(self, partitions: int = 1) -> ContextManager[Topic]:
         raise NotImplementedError
@@ -397,6 +399,11 @@ class StreamsTestMixin(ABC, Generic[TStrategyPayload]):
     def test_pause_resume_rebalancing(self) -> None:
         payloads = self.get_payloads()
 
+        consumer_a_on_assign = mock.Mock()
+        consumer_a_on_revoke = mock.Mock()
+        consumer_b_on_assign = mock.Mock()
+        consumer_b_on_revoke = mock.Mock()
+
         with self.get_topic(2) as topic, closing(
             self.get_producer()
         ) as producer, closing(
@@ -408,10 +415,22 @@ class StreamsTestMixin(ABC, Generic[TStrategyPayload]):
                 producer.produce(Partition(topic, i), next(payloads)).result(
                     timeout=5.0
                 )
-                for i in range(2)
+                for i in [0, 1]
             ]
 
-            consumer_a.subscribe([topic])
+            def wait_until_rebalancing(
+                from_consumer: Consumer[Any], to_consumer: Consumer[Any]
+            ) -> None:
+                for _ in range(10):
+                    assert from_consumer.poll(0) is None
+                    if to_consumer.poll(1.0) is not None:
+                        return
+
+                raise RuntimeError("no rebalancing happened")
+
+            consumer_a.subscribe(
+                [topic], on_assign=consumer_a_on_assign, on_revoke=consumer_a_on_revoke
+            )
 
             # It doesn't really matter which message is fetched first -- we
             # just want to know the assignment occurred.
@@ -428,19 +447,69 @@ class StreamsTestMixin(ABC, Generic[TStrategyPayload]):
                 [Partition(topic, 0), Partition(topic, 1)]
             )
 
-            consumer_b.subscribe([topic])
-            for i in range(10):
-                assert consumer_a.poll(0) is None  # attempt to force session timeout
-                if consumer_b.poll(1.0) is not None:
-                    break
-            else:
-                assert False, "rebalance did not occur"
+            consumer_b.subscribe(
+                [topic], on_assign=consumer_b_on_assign, on_revoke=consumer_b_on_revoke
+            )
+
+            wait_until_rebalancing(consumer_a, consumer_b)
 
-            # The first consumer should have had its offsets rolled back, as
-            # well as should have had it's partition resumed during
-            # rebalancing.
-            assert consumer_a.paused() == []
-            assert consumer_a.poll(10.0) is not None
+            if self.cooperative_sticky:
+                # within incremental rebalancing, only one partition should have been reassigned to the consumer_b, and consumer_a should remain paused
+                assert consumer_a.paused() == [Partition(topic, 1)]
+                assert consumer_a.poll(10.0) is None
+            else:
+                # The first consumer should have had its offsets rolled back, as
+                # well as should have had it's partition resumed during
+                # rebalancing.
+                assert consumer_a.paused() == []
+                assert consumer_a.poll(10.0) is not None
 
             assert len(consumer_a.tell()) == 1
             assert len(consumer_b.tell()) == 1
+
+            (consumer_a_partition,) = consumer_a.tell()
+            (consumer_b_partition,) = consumer_b.tell()
+
+            # Pause consumer_a again.
+            consumer_a.pause(list(consumer_a.tell()))
+            # if we close consumer_a, consumer_b should get all partitions
+            producer.produce(next(iter(consumer_a.tell())), next(payloads)).result(
+                timeout=5.0
+            )
+            consumer_a.unsubscribe()
+            wait_until_rebalancing(consumer_a, consumer_b)
+
+            assert len(consumer_b.tell()) == 2
+
+            if self.cooperative_sticky:
+
+                assert consumer_a_on_assign.mock_calls == [
+                    mock.call({Partition(topic, 0): 0, Partition(topic, 1): 0}),
+                ]
+                assert consumer_a_on_revoke.mock_calls == [
+                    mock.call([Partition(topic, 0)]),
+                    mock.call([Partition(topic, 1)]),
+                ]
+
+                assert consumer_b_on_assign.mock_calls == [
+                    mock.call({Partition(topic, 0): 0}),
+                    mock.call({Partition(topic, 1): 0}),
+                ]
+                assert consumer_b_on_revoke.mock_calls == []
+            else:
+                assert consumer_a_on_assign.mock_calls == [
+                    mock.call({Partition(topic, 0): 0, Partition(topic, 1): 0}),
+                    mock.call({consumer_a_partition: 0}),
+                ]
+                assert consumer_a_on_revoke.mock_calls == [
+                    mock.call([Partition(topic, 0), Partition(topic, 1)]),
+                    mock.call([consumer_a_partition]),
+                ]
+
+                assert consumer_b_on_assign.mock_calls == [
+                    mock.call({consumer_b_partition: 0}),
+                    mock.call({Partition(topic, 0): 0, Partition(topic, 1): 0}),
+                ]
+                assert consumer_b_on_revoke.mock_calls == [
+                    mock.call([consumer_b_partition])
+                ]

tests/backends/test_kafka.py

@@ -14,7 +14,10 @@ from confluent_kafka.admin import AdminClient, NewTopic
 
 from arroyo.backends.kafka import KafkaConsumer, KafkaPayload, KafkaProducer
 from arroyo.backends.kafka.commit import CommitCodec
-from arroyo.backends.kafka.configuration import build_kafka_configuration
+from arroyo.backends.kafka.configuration import (
+    KafkaBrokerConfig,
+    build_kafka_configuration,
+)
 from arroyo.backends.kafka.consumer import as_kafka_configuration_bool
 from arroyo.commit import IMMEDIATE, Commit
 from arroyo.errors import ConsumerError, EndOfPartition
@@ -56,6 +59,7 @@ def get_topic(
     configuration: Mapping[str, Any], partitions_count: int
 ) -> Iterator[Topic]:
     name = f"test-{uuid.uuid1().hex}"
+    configuration = dict(configuration)
     client = AdminClient(configuration)
     [[key, future]] = client.create_topics(
         [NewTopic(name, num_partitions=partitions_count, replication_factor=1)]
@@ -71,10 +75,15 @@ def get_topic(
 
 
 class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
+    kip_848 = False
 
-    configuration = build_kafka_configuration(
-        {"bootstrap.servers": os.environ.get("DEFAULT_BROKERS", "localhost:9092")}
-    )
+    @property
+    def configuration(self) -> KafkaBrokerConfig:
+        config = {
+            "bootstrap.servers": os.environ.get("DEFAULT_BROKERS", "localhost:9092"),
+        }
+
+        return build_kafka_configuration(config)
 
     @contextlib.contextmanager
     def get_topic(self, partitions: int = 1) -> Iterator[Topic]:
@@ -90,7 +99,7 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
         enable_end_of_partition: bool = True,
         auto_offset_reset: str = "earliest",
         strict_offset_reset: Optional[bool] = None,
-        max_poll_interval_ms: Optional[int] = None
+        max_poll_interval_ms: Optional[int] = None,
     ) -> KafkaConsumer:
         configuration = {
             **self.configuration,
@@ -110,6 +119,16 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
             if max_poll_interval_ms < 45000:
                 configuration["session.timeout.ms"] = max_poll_interval_ms
 
+        if self.cooperative_sticky:
+            configuration["partition.assignment.strategy"] = "cooperative-sticky"
+
+        if self.kip_848:
+            configuration["group.protocol"] = "consumer"
+            configuration.pop("session.timeout.ms")
+            configuration.pop("max.poll.interval.ms", None)
+            assert "group.protocol.type" not in configuration
+            assert "heartbeat.interval.ms" not in configuration
+
         return KafkaConsumer(configuration)
 
     def get_producer(self) -> KafkaProducer:
@@ -210,7 +229,9 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
         poll_interval = 6000
 
         with self.get_topic() as topic:
-            with closing(self.get_producer()) as producer, closing(self.get_consumer(max_poll_interval_ms=poll_interval)) as consumer:
+            with closing(self.get_producer()) as producer, closing(
+                self.get_consumer(max_poll_interval_ms=poll_interval)
+            ) as consumer:
                 producer.produce(topic, next(self.get_payloads())).result(5.0)
 
                 processor = StreamProcessor(consumer, topic, factory, IMMEDIATE)
@@ -245,6 +266,16 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
                 assert consumer.paused() == []
 
 
+class TestKafkaStreamsIncrementalRebalancing(TestKafkaStreams):
+    # re-test the kafka consumer with cooperative-sticky rebalancing
+    cooperative_sticky = True
+
+
+@pytest.mark.skip("kip-848 not functional yet")
+class TestKafkaStreamsKip848(TestKafkaStreams):
+    kip_848 = True
+
+
 def test_commit_codec() -> None:
     commit = Commit(
         "group", Partition(Topic("topic"), 0), 0, time.time(), time.time() - 5

tests/backends/test_local.py

@@ -18,14 +18,14 @@ from arroyo.backends.local.storages.abstract import (
 )
 from arroyo.backends.local.storages.memory import MemoryMessageStorage
 from arroyo.types import Partition, Topic
-from arroyo.utils.clock import TestingClock
+from arroyo.utils.clock import MockedClock
 from tests.backends.mixins import StreamsTestMixin
 
 
 class LocalStreamsTestMixin(StreamsTestMixin[int]):
     def setUp(self) -> None:
         self.storage = self.get_message_storage()
-        self.broker = LocalBroker(self.storage, TestingClock())
+        self.broker = LocalBroker(self.storage, MockedClock())
 
     @abstractmethod
     def get_message_storage(self) -> MessageStorage[int]:

tests/processing/strategies/test_produce.py

@@ -9,13 +9,13 @@ from arroyo.backends.local.storages.memory import MemoryMessageStorage
 from arroyo.processing.strategies.abstract import MessageRejected
 from arroyo.processing.strategies.produce import Produce
 from arroyo.types import Message, Partition, Topic, Value
-from arroyo.utils.clock import TestingClock
+from arroyo.utils.clock import MockedClock
 
 
 def test_produce() -> None:
     orig_topic = Topic("orig-topic")
     result_topic = Topic("result-topic")
-    clock = TestingClock()
+    clock = MockedClock()
     broker_storage: MemoryMessageStorage[KafkaPayload] = MemoryMessageStorage()
     broker: LocalBroker[KafkaPayload] = LocalBroker(broker_storage, clock)
     broker.create_topic(result_topic, partitions=1)

tests/processing/test_processor.py

@@ -19,7 +19,7 @@ from arroyo.processing.strategies.abstract import (
     ProcessingStrategyFactory,
 )
 from arroyo.types import BrokerValue, Commit, Message, Partition, Topic
-from arroyo.utils.clock import TestingClock
+from arroyo.utils.clock import MockedClock
 from tests.assertions import assert_changes, assert_does_not_change
 from tests.metrics import Increment, TestingMetricsBackend, Timing
 
@@ -41,6 +41,7 @@ def test_stream_processor_lifecycle() -> None:
 
     # The processor should accept heartbeat messages without an assignment or
     # active processor.
+    consumer.tell.return_value = {}
     consumer.poll.return_value = None
     processor._run_once()
 
@@ -166,6 +167,7 @@ def test_stream_processor_termination_on_error() -> None:
     offset = 0
     now = datetime.now()
 
+    consumer.tell.return_value = {}
     consumer.poll.return_value = BrokerValue(0, partition, offset, now)
 
     exception = NotImplementedError("error")
@@ -199,6 +201,7 @@ def test_stream_processor_invalid_message_from_poll() -> None:
     offset = 1
     now = datetime.now()
 
+    consumer.tell.return_value = {}
     consumer.poll.side_effect = [BrokerValue(0, partition, offset, now)]
 
     strategy = mock.Mock()
@@ -236,6 +239,7 @@ def test_stream_processor_invalid_message_from_submit() -> None:
     offset = 1
     now = datetime.now()
 
+    consumer.tell.return_value = {}
     consumer.poll.side_effect = [
         BrokerValue(0, partition, offset, now),
         BrokerValue(1, partition, offset + 1, now),
@@ -283,6 +287,7 @@ def test_stream_processor_create_with_partitions() -> None:
     topic = Topic("topic")
 
     consumer = mock.Mock()
+    consumer.tell.return_value = {}
     strategy = mock.Mock()
     factory = mock.Mock()
     factory.create_with_partitions.return_value = strategy
@@ -306,13 +311,15 @@ def test_stream_processor_create_with_partitions() -> None:
     assert factory.create_with_partitions.call_count == 1
     assert create_args[1] == offsets_p0
 
+    consumer.tell.return_value = {**offsets_p0}
+
     # Second partition assigned
     offsets_p1 = {Partition(topic, 1): 0}
     assignment_callback(offsets_p1)
 
     create_args, _ = factory.create_with_partitions.call_args
     assert factory.create_with_partitions.call_count == 2
-    assert create_args[1] == offsets_p1
+    assert create_args[1] == {**offsets_p1, **offsets_p0}
 
     processor._run_once()
 
@@ -376,6 +383,7 @@ def run_commit_policy_test(
 ) -> Sequence[int]:
     commit = mock.Mock()
     consumer = mock.Mock()
+    consumer.tell.return_value = {}
     consumer.commit_offsets = commit
 
     factory = CommitOffsetsFactory()
@@ -523,7 +531,7 @@ def test_commit_policy_bench(
     storage: MessageStorage[int] = MemoryMessageStorage()
     storage.create_topic(topic, num_partitions)
 
-    broker = LocalBroker(storage, TestingClock())
+    broker = LocalBroker(storage, MockedClock())
 
     consumer = broker.get_consumer("test-group", enable_end_of_partition=True)
 
@@ -551,6 +559,7 @@ def test_dlq() -> None:
     partition = Partition(topic, 0)
     consumer = mock.Mock()
     consumer.poll.return_value = BrokerValue(0, partition, 1, datetime.now())
+    consumer.tell.return_value = {}
     strategy = mock.Mock()
     strategy.submit.side_effect = InvalidMessage(partition, 1)
     factory = mock.Mock()
@@ -585,6 +594,7 @@ def test_healthcheck(tmpdir: py.path.local) -> None:
     consumer = mock.Mock()
     now = datetime.now()
     consumer.poll.return_value = BrokerValue(0, partition, 1, now)
+    consumer.tell.return_value = {}
     strategy = mock.Mock()
     strategy.submit.side_effect = InvalidMessage(partition, 1)
     factory = mock.Mock()

tests/utils/test_retries.py

@@ -3,7 +3,7 @@ from unittest import mock
 
 import pytest
 
-from arroyo.utils.clock import TestingClock
+from arroyo.utils.clock import MockedClock
 from arroyo.utils.retries import BasicRetryPolicy, RetryException, constant_delay
 
 value = object()
@@ -45,7 +45,7 @@ def setup_function() -> None:
 
 def test_basic_retry_policy_no_delay() -> None:
 
-    clock = TestingClock()
+    clock = MockedClock()
 
     policy = BasicRetryPolicy(3, clock=clock)
 
@@ -68,19 +68,19 @@ def test_basic_retry_policy_no_delay() -> None:
 
 @pytest.mark.parametrize("delay", [1, constant_delay(1)])
 def test_basic_retry_policy_with_delay(delay: int) -> None:
-    clock = TestingClock()
+    clock = MockedClock()
     policy = BasicRetryPolicy(3, delay, clock=clock)
     assert policy.call(good_function) is value
     assert good_function.call_count == 1
     assert clock.time() == 0
 
-    clock = TestingClock()
+    clock = MockedClock()
     policy = BasicRetryPolicy(3, delay, clock=clock)
     assert policy.call(flaky_function) is value
     assert flaky_function.call_count == 2
     assert clock.time() == 1  # one retry
 
-    clock = TestingClock()
+    clock = MockedClock()
     policy = BasicRetryPolicy(3, delay, clock=clock)
     try:
         policy.call(bad_function)
@@ -109,7 +109,7 @@ def test_basic_retry_policy_with_supression() -> None:
     def suppression_test(exception: Exception) -> bool:
         return isinstance(exception, ExpectedError)
 
-    clock = TestingClock()
+    clock = MockedClock()
     policy = BasicRetryPolicy(
         3, constant_delay(1), suppression_test=suppression_test, clock=clock
     )