mirror of
https://github.com/Sea-Haven-Industries/afterhours-shift-manager.git
synced 2026-09-30 07:53:11 +00:00
65 lines
2 KiB
Python
65 lines
2 KiB
Python
|
"""copy_dynamodb retries UnprocessedItems instead of counting them as written."""
|
||
|
|
|
||
|
|
import importlib.util
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
spec = importlib.util.spec_from_file_location(
|
||
|
|
"copy_dynamodb", ROOT / "scripts" / "cutover" / "copy_dynamodb.py"
|
||
|
|
)
|
||
|
|
mod = importlib.util.module_from_spec(spec)
|
||
|
|
sys.modules["copy_dynamodb"] = mod
|
||
|
|
spec.loader.exec_module(mod)
|
||
|
|
return mod
|
||
|
|
|
||
|
|
|
||
|
|
mod = _load()
|
||
|
|
|
||
|
|
|
||
|
|
class FakeDdb:
|
||
|
|
def __init__(self, unprocessed_first=None):
|
||
|
|
self.calls = []
|
||
|
|
self.unprocessed_first = list(unprocessed_first or [])
|
||
|
|
self._first = True
|
||
|
|
|
||
|
|
def batch_write_item(self, RequestItems):
|
||
|
|
batch = RequestItems[mod.TABLE]
|
||
|
|
self.calls.append(batch)
|
||
|
|
if self._first and self.unprocessed_first:
|
||
|
|
self._first = False
|
||
|
|
leftover = [req for req in batch if req in self.unprocessed_first]
|
||
|
|
return {"UnprocessedItems": {mod.TABLE: leftover} if leftover else {}}
|
||
|
|
return {"UnprocessedItems": {}}
|
||
|
|
|
||
|
|
|
||
|
|
def test_batch_write_retries_unprocessed_items():
|
||
|
|
items = [{"PK": {"S": "a"}}, {"PK": {"S": "b"}}]
|
||
|
|
first = [{"PutRequest": {"Item": items[1]}}]
|
||
|
|
client = FakeDdb(unprocessed_first=first)
|
||
|
|
written = mod._batch_write_all(client, items, sleep=lambda _s: None)
|
||
|
|
assert written == 2
|
||
|
|
assert len(client.calls) == 2
|
||
|
|
assert client.calls[1] == first
|
||
|
|
|
||
|
|
|
||
|
|
def test_batch_write_raises_if_unprocessed_remain():
|
||
|
|
items = [{"PK": {"S": "a"}}]
|
||
|
|
stuck = [{"PutRequest": {"Item": items[0]}}]
|
||
|
|
client = FakeDdb(unprocessed_first=stuck)
|
||
|
|
client._always = True
|
||
|
|
|
||
|
|
def always_unprocessed(RequestItems):
|
||
|
|
client.calls.append(RequestItems[mod.TABLE])
|
||
|
|
return {"UnprocessedItems": {mod.TABLE: stuck}}
|
||
|
|
|
||
|
|
client.batch_write_item = always_unprocessed
|
||
|
|
try:
|
||
|
|
mod._batch_write_all(client, items, sleep=lambda _s: None, max_attempts=3)
|
||
|
|
raise AssertionError("expected RuntimeError")
|
||
|
|
except RuntimeError as exc:
|
||
|
|
assert "unprocessed" in str(exc)
|