diff --git a/json_loader.py b/json_loader.py index d158892..ecc9487 100644 --- a/json_loader.py +++ b/json_loader.py @@ -21,11 +21,15 @@ def to_int(val: Any) -> int: return 0 def get_batch_item(data: dict[str, Any], sequence_number: int) -> dict[str, Any]: - """Resolve batch item by sequence_number, clamping to valid range.""" + """Resolve batch item by sequence_number field, falling back to array index.""" if KEY_BATCH_DATA in data and isinstance(data[KEY_BATCH_DATA], list) and len(data[KEY_BATCH_DATA]) > 0: + # Search by sequence_number field first + for item in data[KEY_BATCH_DATA]: + if int(item.get("sequence_number", 0)) == sequence_number: + return item + # Fallback to array index idx = max(0, min(sequence_number - 1, len(data[KEY_BATCH_DATA]) - 1)) - if sequence_number - 1 != idx: - logger.warning(f"Sequence {sequence_number} out of range (1-{len(data[KEY_BATCH_DATA])}), clamped to {idx + 1}") + logger.warning(f"No item with sequence_number={sequence_number}, falling back to index {idx}") return data[KEY_BATCH_DATA][idx] return data diff --git a/tests/test_json_loader.py b/tests/test_json_loader.py index 0f9b3e4..2a02690 100644 --- a/tests/test_json_loader.py +++ b/tests/test_json_loader.py @@ -31,7 +31,16 @@ class TestToInt: class TestGetBatchItem: - def test_valid_index(self): + def test_lookup_by_sequence_number_field(self): + data = {"batch_data": [ + {"sequence_number": 1, "a": "first"}, + {"sequence_number": 5, "a": "fifth"}, + {"sequence_number": 3, "a": "third"}, + ]} + assert get_batch_item(data, 5) == {"sequence_number": 5, "a": "fifth"} + assert get_batch_item(data, 3) == {"sequence_number": 3, "a": "third"} + + def test_fallback_to_index(self): data = {"batch_data": [{"a": 1}, {"a": 2}, {"a": 3}]} assert get_batch_item(data, 2) == {"a": 2}