test: Fix flaky test

arroyo/backends/abstract.py

@@ -139,9 +139,7 @@ def seek(self, offsets: Mapping[Partition, int]) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def stage_positions(
-        self, positions: Mapping[Partition, Position], force: bool = False
-    ) -> None:
+    def stage_positions(self, positions: Mapping[Partition, Position]) -> None:
         """
         Stage offsets to be committed. If an offset has already been staged
         for a given partition, that offset is overwritten (even if the offset

arroyo/backends/kafka/consumer.py

@@ -582,22 +582,16 @@ def paused(self) -> Sequence[Partition]:
 
         return [*self.__paused]
 
-    def stage_positions(
-        self, positions: Mapping[Partition, Position], force: bool = False
-    ) -> None:
+    def stage_positions(self, positions: Mapping[Partition, Position]) -> None:
         if self.__state in {KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR}:
             raise InvalidState(self.__state)
 
-        if not force:
-            if positions.keys() - self.__offsets.keys():
-                raise ConsumerError("cannot stage offsets for unassigned partitions")
+        if positions.keys() - self.__offsets.keys():
+            raise ConsumerError("cannot stage offsets for unassigned partitions")
 
-            self.__validate_offsets(
-                {
-                    partition: position.offset
-                    for (partition, position) in positions.items()
-                }
-            )
+        self.__validate_offsets(
+            {partition: position.offset for (partition, position) in positions.items()}
+        )
 
         # TODO: Maybe log a warning if these offsets exceed the current
         # offsets, since that's probably a side effect of an incorrect usage

arroyo/backends/local/backend.py

@@ -304,25 +304,20 @@ def seek(self, offsets: Mapping[Partition, int]) -> None:
 
             self.__offsets.update(offsets)
 
-    def stage_positions(
-        self, positions: Mapping[Partition, Position], force: bool = False
-    ) -> None:
+    def stage_positions(self, positions: Mapping[Partition, Position]) -> None:
         with self.__lock:
             if self.__closed:
                 raise RuntimeError("consumer is closed")
 
-            if not force:
-                if positions.keys() - self.__offsets.keys():
-                    raise ConsumerError(
-                        "cannot stage offsets for unassigned partitions"
-                    )
+            if positions.keys() - self.__offsets.keys():
+                raise ConsumerError("cannot stage offsets for unassigned partitions")
 
-                self.__validate_offsets(
-                    {
-                        partition: position.offset
-                        for (partition, position) in positions.items()
-                    }
-                )
+            self.__validate_offsets(
+                {
+                    partition: position.offset
+                    for (partition, position) in positions.items()
+                }
+            )
 
             self.__staged_positions.update(positions)
 

arroyo/synchronized.py

@@ -1,8 +1,8 @@
 import logging
-from datetime import datetime
-from time import time
 from dataclasses import dataclass
+from datetime import datetime
 from threading import Event
+from time import time
 from typing import Callable, Mapping, MutableMapping, Optional, Sequence, Set
 
 from arroyo.backends.abstract import Consumer
@@ -315,10 +315,8 @@ def tell(self) -> Mapping[Partition, int]:
     def seek(self, offsets: Mapping[Partition, int]) -> None:
         return self.__consumer.seek(offsets)
 
-    def stage_positions(
-        self, positions: Mapping[Partition, Position], force: bool = False
-    ) -> None:
-        return self.__consumer.stage_positions(positions, force=force)
+    def stage_positions(self, positions: Mapping[Partition, Position]) -> None:
+        return self.__consumer.stage_positions(positions)
 
     def commit_positions(self) -> Mapping[Partition, Position]:
         return self.__consumer.commit_positions()

tests/backends/test_kafka.py

@@ -17,7 +17,7 @@
 from arroyo.backends.kafka.consumer import as_kafka_configuration_bool
 from arroyo.errors import ConsumerError, EndOfPartition
 from arroyo.synchronized import Commit, commit_codec
-from arroyo.types import Message, Partition, Topic, Position
+from arroyo.types import Message, Partition, Position, Topic
 from tests.backends.mixins import StreamsTestMixin
 
 
@@ -143,13 +143,13 @@ def test_lenient_offset_reset_latest(self) -> None:
             # write a nonsense offset
             with closing(self.get_consumer(strict_offset_reset=False)) as consumer:
                 consumer.subscribe([topic])
+                consumer.poll(10.0)  # Wait for assignment
                 consumer.stage_positions(
                     {
                         message.partition: Position(
                             offset=message.offset + 1000, timestamp=message.timestamp
                         )
                     },
-                    force=True,
                 )
                 consumer.commit_positions()