Fix setitem loop now that we use Dynamo retry config. Update number f…

…or retries in test assertion
Yelp · Jan 20, 2025 · f88d7f9 · f88d7f9
1 parent 3e74d75
commit f88d7f9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tests/serialize/runstate/dynamodb_state_store_test.py b/tests/serialize/runstate/dynamodb_state_store_test.py
@@ -345,7 +345,7 @@ def test_retry_reading(self, store, small_object, large_object):
                 ):
                     store.restore(keys)
             except Exception:
-                assert_equal(mock_failed_read.call_count, 11)
+                assert_equal(mock_failed_read.call_count, 10)
 
     def test_restore_exception_propagation(self, store, small_object):
         # This test is to ensure that restore propagates exceptions upwards: see DAR-2328

diff --git a/tron/serialize/runstate/dynamodb_state_store.py b/tron/serialize/runstate/dynamodb_state_store.py
@@ -374,7 +374,9 @@ def __setitem__(self, key: str, value: Tuple[bytes, str]) -> None:
 
             items.append(item)
 
-            while len(items) == MAX_TRANSACT_WRITE_ITEMS or index == max_partitions - 1:
+            # We want to write the items when we've either reached the max number of items
+            # for a transaction, or when we're done processing all partitions
+            if len(items) == MAX_TRANSACT_WRITE_ITEMS or index == max_partitions - 1:
                 try:
                     self.client.transact_write_items(TransactItems=items)
                     items = []