diff --git a/sagemaker-core/src/sagemaker/core/iterators.py b/sagemaker-core/src/sagemaker/core/iterators.py index 049a3a3cf5..abdaab4e05 100644 --- a/sagemaker-core/src/sagemaker/core/iterators.py +++ b/sagemaker-core/src/sagemaker/core/iterators.py @@ -108,7 +108,7 @@ def __next__(self): # handle API response errors and force terminate. handle_stream_errors(chunk) # print and move on to next response byte - print("Unknown event type:" + chunk) + print("Unknown event type:", chunk) continue return chunk["PayloadPart"]["Bytes"] @@ -190,7 +190,7 @@ def __next__(self): # handle API response errors and force terminate. handle_stream_errors(chunk) # print and move on to next response byte - print("Unknown event type:" + chunk) + print("Unknown event type:", chunk) continue # Check buffer size before writing to prevent unbounded memory consumption diff --git a/sagemaker-core/tests/unit/test_iterators.py b/sagemaker-core/tests/unit/test_iterators.py index d9c6831ce9..fcefbcb3d5 100644 --- a/sagemaker-core/tests/unit/test_iterators.py +++ b/sagemaker-core/tests/unit/test_iterators.py @@ -254,3 +254,19 @@ def test_line_iterator_no_trailing_newline_at_end(): with pytest.raises(StopIteration): next(iterator) + + +def test_byte_iterator_skips_unknown_event_type(): + """Test that ByteIterator skips events it does not recognize.""" + event_stream = [{"SomeNewEvent": {}}, {"PayloadPart": {"Bytes": b"hello"}}] + iterator = ByteIterator(event_stream) + + assert next(iterator) == b"hello" + + +def test_line_iterator_skips_unknown_event_type(): + """Test that LineIterator skips events it does not recognize.""" + event_stream = [{"SomeNewEvent": {}}, {"PayloadPart": {"Bytes": b"hello\n"}}] + iterator = LineIterator(event_stream) + + assert next(iterator) == b"hello"